m-thesis-introduction/simulations/airshower_beacon_simulation/lib/snr.py

56 lines
1.4 KiB
Python

import numpy as np
from collections import namedtuple
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):
fft, freqs = get_freq_spec(samples, samplerate)
bandmask = bandpass_mask(freqs, band=band)
if normalise_bandsize:
bins = np.count_nonzero(bandmask, axis=-1)
else:
bins = 1
power = np.sum(np.abs(fft[bandmask])**2)
return power/bins
def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None):
if noise_band is None:
noise_band = signal_band
if noise is None:
noise = samples
noise_power = bandpower(noise, samplerate, noise_band)
signal_power = bandpower(samples, samplerate, signal_band)
return (signal_power/noise_power)**0.5