57 lines
2.2 KiB
Python
57 lines
2.2 KiB
Python
import numpy as np
|
|
|
|
|
|
def gauss(N, mu, sigma):
|
|
x = np.arange(N)
|
|
norm = sigma * np.sqrt(2*np.pi)
|
|
return np.exp(-1/2 * (x - mu)**2 / sigma**2) / norm
|
|
|
|
|
|
def shift(N, n, x):
|
|
"""shift the center of the signal 'x' to time index 'n'."""
|
|
y = np.zeros(N)
|
|
xl = x.shape[0]
|
|
s = n - xl // 2
|
|
x_bi, x_ei = np.clip(xl // 2 - n, a_min=0, a_max=xl), np.clip(N + xl - xl // 2 - n - xl % 2, a_min=0, a_max=xl)
|
|
y_bi, y_ei = np.clip(n - xl // 2, a_min=0, a_max=N), np.clip(n + xl // 2 + xl % 2, a_min=0, a_max=N)
|
|
y[y_bi:y_ei] = x[x_bi:x_ei]
|
|
return y
|
|
|
|
|
|
class SigQuality:
|
|
"""
|
|
Compute the Signal-to-Noise Ratio of beats
|
|
"""
|
|
gauss_beat_template_win_sec = 0.25542 #: gauss window width (as compared to beats in ssf function)
|
|
gauss_beat_template_sigma_sec = 0.027 #: gauss bump half-width parameter (as compared to beats in ssf function)
|
|
#gauss_amplitude = 2.0
|
|
|
|
def __init__(self): pass
|
|
|
|
def get_snr(self, fs, ssf, ssf_threshold, ssf_zxings):
|
|
"""Compute the Signal-to-Noise Ratio of beats, based on SSF function and detected beat locations."""
|
|
sigma = fs * self.gauss_beat_template_sigma_sec
|
|
W = int(fs * self.gauss_beat_template_win_sec)
|
|
gb = gauss(W, W//2, sigma)
|
|
# place gaussians on estimated beat locations
|
|
ssf_est = np.zeros(ssf.shape[0])
|
|
for i in ssf_zxings:
|
|
ssf_est += shift(ssf.shape[0], i, gb)
|
|
ssf_est /= gb[W//2] # normalize amplitude to 1.0
|
|
ssf_est = np.roll(ssf_est, int(sigma)) # shift to right (beat loc = gauss beginning, not center)
|
|
#sqi_ref = ssf_est * (self.gauss_amplitude * ssf_threshold) # set amplitude
|
|
|
|
# penalty term = (where there should be no amplitude, according to gaussians)
|
|
sqi_pen = 1.0 - ssf_est
|
|
|
|
# sqi = ratio of signal energy in goal vs. in noise
|
|
sqi_goal = np.sum(ssf_est * (ssf**2))
|
|
sqi_noise = np.sum(sqi_pen * (ssf**2))
|
|
|
|
# noise is everywhere, while signal is only around detected peaks - correct for this.
|
|
goal_density = np.mean(np.clip(2*sigma / np.diff(ssf_zxings), a_min=0, a_max=1))
|
|
sqi_goal /= goal_density
|
|
sqi = 10 * (np.log10(sqi_goal) - np.log10(sqi_noise))
|
|
|
|
return sqi
|