improved heart rate calculations
This commit is contained in:
@@ -4526,9 +4526,9 @@ def short_channel_processing_for_hr(
|
||||
|
||||
|
||||
|
||||
def calculate_heart_rate_neurokit(sfreq: float, signal_trimmed: NDArray[float64], hr_low_freq, hr_high_freq, max_low_hr, max_high_hr, smoothing_window_hr) -> tuple[NDArray[float64], float]:
|
||||
def calculate_heart_rate_neurokit(sfreq: float, signal_trimmed: NDArray[float64], hr_low_freq, hr_high_freq, max_low_hr, max_high_hr, smoothing_window_hr, short_channels) -> tuple[NDArray[float64], float]:
|
||||
"""
|
||||
Calculate and smooth heart rate from a trimmed signal using NeuroKit.
|
||||
Calculate and smooths heart rate from a trimmed signal using NeuroKit.
|
||||
|
||||
Parameters
|
||||
----------
|
||||
@@ -4543,39 +4543,50 @@ def calculate_heart_rate_neurokit(sfreq: float, signal_trimmed: NDArray[float64]
|
||||
- NDArray[float64]: Smoothed heart rate time series (BPM).
|
||||
- float: Mean heart rate.
|
||||
"""
|
||||
|
||||
logger.info("Calculating heart rate using NeuroKit...")
|
||||
|
||||
# Filter signal to isolate heart rate frequencies and detect peaks
|
||||
logger.info("Filtering the signal and detecting peaks...")
|
||||
signal_filtered = cast(NDArray[float64], nk.signal_filter(signal_trimmed, sampling_rate=sfreq, lowcut=hr_low_freq, highcut=hr_high_freq)) # type: ignore
|
||||
peaks_dict = cast(dict[str, Any], nk.signal_findpeaks(signal_filtered)) # type: ignore
|
||||
peaks = peaks_dict['Peaks']
|
||||
hr = cast(NDArray[float64], nk.signal_rate(peaks, sampling_rate=sfreq, desired_length=len(signal_trimmed))) # type: ignore
|
||||
logger.info("Filtering the signal and detecting pulsatile peaks...")
|
||||
signal_filtered = cast(NDArray[float64], nk.signal_filter(signal_trimmed, sampling_rate=sfreq, lowcut=hr_low_freq, highcut=hr_high_freq))
|
||||
|
||||
# bishop works better with challenging datasets, but like it is super slow on good datasets?
|
||||
if short_channels:
|
||||
peaks_dict = cast(dict[str, Any], nk.ppg_findpeaks(signal_filtered, sampling_rate=sfreq, method="elgendi"))
|
||||
else:
|
||||
peaks_dict = cast(dict[str, Any], nk.ppg_findpeaks(signal_filtered, sampling_rate=sfreq, method="bishop"))
|
||||
peaks = peaks_dict['PPG_Peaks']
|
||||
logger.info(f"ppg_findpeaks found {len(peaks)} peaks over {len(signal_trimmed)/sfreq:.1f}s "
|
||||
f"(~{len(peaks) / (len(signal_trimmed)/sfreq) * 60:.1f} BPM implied by peak count alone)")
|
||||
|
||||
if len(peaks) < 2:
|
||||
logger.warning("ppg_findpeaks found fewer than 2 peaks - heart rate estimate is unreliable.")
|
||||
return np.full(len(signal_trimmed), np.nan), float('nan')
|
||||
|
||||
hr = cast(NDArray[float64], nk.signal_rate(peaks, sampling_rate=sfreq, desired_length=len(signal_trimmed)))
|
||||
logger.info(f"Pre-clip HR range: min={hr.min():.1f}, max={hr.max():.1f} BPM (max_low_hr={max_low_hr}, max_high_hr={max_high_hr})")
|
||||
|
||||
hr_clean = np.clip(hr, max_low_hr, max_high_hr)
|
||||
|
||||
# Smooth heart rate time series by replacing spikes with local rolling mean and calculate the mean
|
||||
logger.info("Smoothing the signal and calculating the mean...")
|
||||
hr_series = pd.Series(hr_clean)
|
||||
local_median = hr_series.rolling(window=smoothing_window_hr, center=True, min_periods=1).median()
|
||||
spikes = hr_series > (local_median + 10)
|
||||
spikes = (hr_series > local_median + 10) | (hr_series < local_median - 10) # was upward-only; catches drops too
|
||||
smoothed_values = hr_series.copy()
|
||||
smoothed_spikes = hr_series.rolling(window=smoothing_window_hr, center=True, min_periods=1).mean()
|
||||
smoothed_values[spikes] = smoothed_spikes[spikes]
|
||||
hr_smooth_nk = cast(NDArray[float64], smoothed_values.to_numpy()) # type: ignore
|
||||
hr_smooth_nk = cast(NDArray[float64], smoothed_values.to_numpy())
|
||||
mean_hr_nk = hr_smooth_nk.mean()
|
||||
|
||||
logger.info("Original HR min/max: %f, %f", hr_clean.min(), hr_clean.max())
|
||||
logger.info("Smoothed HR min/max:%f, %f", hr_smooth_nk.min(), hr_smooth_nk.max())
|
||||
logger.info(f"Estimated mean HR nk: {mean_hr_nk:.1f} BPM")
|
||||
|
||||
logger.info("Successfully calculated heart rate using NeuroKit.")
|
||||
|
||||
return hr_smooth_nk, mean_hr_nk
|
||||
|
||||
|
||||
|
||||
def calculate_heart_rate_scipy(sfreq: float, signal_trimmed: NDArray[float64], search_min, search_max) -> tuple[NDArray[floating[Any]], NDArray[float64], np.ndarray[Any, np.dtype[np.bool_]], float]:
|
||||
def calculate_heart_rate_scipy(
|
||||
sfreq: float, signal_trimmed: NDArray[float64], search_min, search_max,
|
||||
cluster_window_bpm: float = 6.0,
|
||||
) -> tuple[NDArray[floating[Any]], NDArray[float64], np.ndarray[Any, np.dtype[np.bool_]], float, float]:
|
||||
"""
|
||||
Estimate heart rate using spectral analysis on a high-pass filtered signal.
|
||||
|
||||
@@ -4585,6 +4596,13 @@ def calculate_heart_rate_scipy(sfreq: float, signal_trimmed: NDArray[float64], s
|
||||
Sampling frequency of the input signal.
|
||||
signal_trimmed : NDArray[float64]
|
||||
Trimmed fNIRS signal to analyze.
|
||||
cluster_window_bpm : float, default 6.0
|
||||
Width (in BPM, +/- from the argmax) used to find nearby local peaks
|
||||
that likely belong to the same underlying cardiac frequency (spread
|
||||
by natural heart-rate variability/frequency modulation) rather than
|
||||
being genuinely separate candidates. The reported HR is the
|
||||
power-weighted centroid of all local peaks within this window of
|
||||
the strongest bin, not just the single tallest bin.
|
||||
|
||||
Returns
|
||||
-------
|
||||
@@ -4594,33 +4612,63 @@ def calculate_heart_rate_scipy(sfreq: float, signal_trimmed: NDArray[float64], s
|
||||
- np.ndarray[Any, np.dtype[np.bool_]]: Boolean mask indicating frequencies within heart rate range.
|
||||
- float: Estimated mean heart rate in BPM corresponding to the PSD peak within the range.
|
||||
"""
|
||||
|
||||
logger.info("Calculating heart rate using SciPy...")
|
||||
|
||||
# Apply a high-pass Butterworth filter to remove slow trends below 0.5 Hz from the trimmed signal (actual data)
|
||||
logger.info("Applying a butterworth filter...")
|
||||
b, a = cast(tuple[NDArray[float64], NDArray[float64]], butter(2, 0.5 / (sfreq / 2), btype='high'))
|
||||
signal_hp = cast(NDArray[float64],filtfilt(b, a, signal_trimmed))
|
||||
signal_hp = cast(NDArray[float64], filtfilt(b, a, signal_trimmed))
|
||||
|
||||
# Calculate the Power Spectral Density (PSD) of the filtered signal using Welch's method
|
||||
logger.info("Calculating the PSD...")
|
||||
nperseg = min(len(signal_hp), 4096)
|
||||
frequencies_scipy, psd_scipy = cast(tuple[NDArray[float64], NDArray[float64]], welch(signal_hp, fs=sfreq, nperseg=nperseg, noverlap=nperseg//2))
|
||||
frequencies_scipy, psd_scipy = cast(tuple[NDArray[float64], NDArray[float64]], welch(signal_hp, fs=sfreq, nperseg=nperseg, noverlap=nperseg // 2))
|
||||
|
||||
# Convert frequency values to beats per minute (BPM) and set a heart rate range
|
||||
logger.info("Converting to BPM...")
|
||||
freq_bpm_scipy = frequencies_scipy * 60
|
||||
freq_range_scipy = (freq_bpm_scipy > search_min) & (freq_bpm_scipy < search_max)
|
||||
|
||||
# Identify the peak frequency within the heart rate range and estimate the mean heart rate in BPM
|
||||
logger.info("Finding the mean...")
|
||||
peak_index = np.argmax(psd_scipy[freq_range_scipy])
|
||||
mean_hr_scipy = freq_bpm_scipy[freq_range_scipy][peak_index]
|
||||
band_bpm = freq_bpm_scipy[freq_range_scipy]
|
||||
band_psd = psd_scipy[freq_range_scipy]
|
||||
if len(band_psd) == 0:
|
||||
raise ValueError(f"No frequency bins fall within the search range ({search_min}-{search_max} BPM).")
|
||||
|
||||
# Find ALL local peaks in the band (bins strictly greater than both neighbors)
|
||||
local_peak_mask = np.zeros(len(band_psd), dtype=bool)
|
||||
if len(band_psd) >= 3:
|
||||
local_peak_mask[1:-1] = (band_psd[1:-1] > band_psd[:-2]) & (band_psd[1:-1] > band_psd[2:])
|
||||
# Edge bins can't be evaluated as local peaks by this rule; if the true
|
||||
# peak sits at the very edge of the search range, argmax below still
|
||||
# catches it as a fallback.
|
||||
peak_indices = np.where(local_peak_mask)[0]
|
||||
if len(peak_indices) == 0:
|
||||
peak_indices = np.array([np.argmax(band_psd)])
|
||||
|
||||
strongest_idx = peak_indices[np.argmax(band_psd[peak_indices])]
|
||||
strongest_bpm = band_bpm[strongest_idx]
|
||||
|
||||
# Cluster: local peaks within cluster_window_bpm of the strongest one
|
||||
cluster_mask = np.abs(band_bpm[peak_indices] - strongest_bpm) <= cluster_window_bpm
|
||||
cluster_indices = peak_indices[cluster_mask]
|
||||
cluster_bpm = band_bpm[cluster_indices]
|
||||
cluster_power = band_psd[cluster_indices]
|
||||
|
||||
# Power-weighted centroid across the cluster - this is the actual fix:
|
||||
# a tight group of near-equal peaks now contributes to ONE combined
|
||||
# estimate near their shared center, instead of a coin-flip winner-take-all.
|
||||
mean_hr_scipy = float(np.average(cluster_bpm, weights=cluster_power))
|
||||
|
||||
# Confidence: strongest cluster's TOTAL power vs. the median power of
|
||||
# everything OUTSIDE the cluster - reflects how dominant the whole
|
||||
# cluster is, not just one bin within it. A 4-peak near-tie spread
|
||||
# across the band now scores lower confidence than a single sharp,
|
||||
# isolated peak of similar height, even though argmax alone couldn't
|
||||
# tell them apart.
|
||||
outside_cluster = np.setdiff1d(np.arange(len(band_psd)), cluster_indices)
|
||||
baseline_power = np.median(band_psd[outside_cluster]) if len(outside_cluster) > 0 else np.median(band_psd)
|
||||
cluster_total_power = cluster_power.sum()
|
||||
peak_confidence = cluster_total_power / baseline_power if baseline_power > 0 else 0.0
|
||||
|
||||
logger.info(f"PSD: {len(cluster_indices)} peak(s) in cluster near {strongest_bpm:.1f} BPM, "
|
||||
f"centroid={mean_hr_scipy:.1f} BPM, confidence={peak_confidence:.2f}x baseline")
|
||||
logger.info("Successfully calculated heart rate using SciPy.")
|
||||
|
||||
return freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy
|
||||
|
||||
return freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy, peak_confidence
|
||||
|
||||
|
||||
def plot_heart_rate(
|
||||
@@ -4896,23 +4944,30 @@ def find_bad_channels_coeff_var(raw, coeff_var_threshold=25.0):
|
||||
|
||||
|
||||
|
||||
def hr_calc(raw, seconds_to_strip_hr, l_freq, h_freq, search_min, search_max, max_low_hr, max_high_hr, smoothing_window_hr, hr_window, short_channels, short_channels_threshold, verbosity):
|
||||
def hr_calc(raw, seconds_to_strip_hr, l_freq, h_freq, search_min, search_max, max_low_hr, max_high_hr, smoothing_window_hr, hr_window, short_channels, short_channels_threshold, verbosity, psd_confidence_threshold: float = 3.0):
|
||||
if short_channels:
|
||||
short_chans = get_short_channels(raw, max_dist=short_channels_threshold)
|
||||
else:
|
||||
short_chans = None
|
||||
sfreq, signal_trimmed, times_trimmed = short_channel_processing_for_hr(raw, short_chans, seconds_to_strip_hr=seconds_to_strip_hr, verbosity=verbosity)
|
||||
hr_smooth_nk, mean_hr_nk = calculate_heart_rate_neurokit(sfreq, signal_trimmed, hr_low_freq=l_freq, hr_high_freq=h_freq, max_low_hr=max_low_hr, max_high_hr=max_high_hr, smoothing_window_hr=smoothing_window_hr)
|
||||
freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy = calculate_heart_rate_scipy(sfreq, signal_trimmed, search_min=search_min, search_max=search_max)
|
||||
hr_smooth_nk, mean_hr_nk = calculate_heart_rate_neurokit(sfreq, signal_trimmed, hr_low_freq=l_freq, hr_high_freq=h_freq, max_low_hr=max_low_hr, max_high_hr=max_high_hr, smoothing_window_hr=smoothing_window_hr, short_channels=short_channels)
|
||||
freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy, psd_confidence = calculate_heart_rate_scipy(sfreq, signal_trimmed, search_min=search_min, search_max=search_max)
|
||||
|
||||
# HACK: This sucks but looking at the graphs I trust neurokit2 more
|
||||
overruled = False
|
||||
if mean_hr_scipy < mean_hr_nk - 15:
|
||||
mean_hr_scipy = mean_hr_nk
|
||||
overruled = True
|
||||
if mean_hr_scipy > mean_hr_nk + 15:
|
||||
mean_hr_scipy = mean_hr_nk
|
||||
overruled = True
|
||||
disagreement = abs(mean_hr_scipy - mean_hr_nk) > 15
|
||||
|
||||
if disagreement:
|
||||
if psd_confidence >= psd_confidence_threshold:
|
||||
logger.info(f"HR estimates disagree ({mean_hr_scipy:.1f} vs {mean_hr_nk:.1f} BPM) - "
|
||||
f"PSD peak is clear (confidence={psd_confidence:.2f}), trusting PSD. Overruling NeuroKit.")
|
||||
mean_hr_nk = mean_hr_scipy
|
||||
overruled = True
|
||||
else:
|
||||
logger.info(f"HR estimates disagree ({mean_hr_scipy:.1f} vs {mean_hr_nk:.1f} BPM) - "
|
||||
f"PSD peak is ambiguous (confidence={psd_confidence:.2f} < {psd_confidence_threshold}), "
|
||||
f"trusting NeuroKit instead. Overruling PSD.")
|
||||
mean_hr_scipy = mean_hr_nk
|
||||
overruled = True
|
||||
|
||||
hr1, hr2 = plot_heart_rate(freq_bpm_scipy, psd_scipy, freq_range_scipy, mean_hr_scipy, hr_smooth_nk, mean_hr_nk, times_trimmed, overruled, hr_window=hr_window)
|
||||
|
||||
|
||||
Reference in New Issue
Block a user