ZH:lib/snr optional debugging plot in function

This commit is contained in:
Eric-Teunis de Boone 2023-02-02 08:49:05 +01:00
parent 91016be038
commit 370b6f366a
2 changed files with 77 additions and 16 deletions

View file

@ -1,6 +1,10 @@
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):
@ -27,30 +31,67 @@ def bandpass_mask(freqs, band=passband()):
return low_pass & high_pass
def bandpower(samples, samplerate=1, band=passband(), normalise_bandsize=True):
fft, freqs = get_freq_spec(samples, samplerate)
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)
bandmask = bandpass_mask(freqs, band=band)
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)
if normalise_bandsize:
bins = np.count_nonzero(bandmask, axis=-1)
real, imag = dtft(band[0], time, samples)
power = np.sum(np.abs(real**2 + imag**2))
else:
bins = 1
bandmask = bandpass_mask(freqs, band=band)
power = np.sum(np.abs(fft[bandmask])**2)
if normalise_bandsize:
bins = np.count_nonzero(bandmask, axis=-1)
else:
bins = 1
return power/bins
bins = max(1, bins)
def signal_to_noise(samples, noise, samplerate=1, signal_band=passband(), noise_band=None):
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):
if noise_band is None:
noise_band = signal_band
if noise is None:
noise = samples
noise_power = bandpower(noise, samplerate, noise_band)
if debug_ax is True:
debug_ax = plt.gca()
signal_power = bandpower(samples, samplerate, signal_band)
noise_power = bandpower(noise, samplerate, noise_band, debug_ax=debug_ax)
signal_power = bandpower(samples, samplerate, signal_band, debug_ax=debug_ax)
return (signal_power/noise_power)**0.5