Extend calculated features to support time, frequency and non_linear domain

rpeaks2hrv.py

@@ -7,24 +7,59 @@ class WindowingMethod(Enum):
     FIRST_INTERVAL = 'first_interval'
     LAST_INTERVAL = 'last_interval'
 
+class FeatureDomain(Enum):
+    TIME = 'time'
+    FREQUENCY = 'freq'
+    NON_LINEAR = 'non_lin'
+
+
+FEATURE_COLS = {
+    FeatureDomain.TIME: [
+        'HRV_MeanNN', 'HRV_SDNN', 'HRV_SDANN1', 'HRV_SDNNI1', 
+        'HRV_SDANN2', 'HRV_SDNNI2', 'HRV_SDANN5', 'HRV_SDNNI5', 
+        'HRV_RMSSD', 'HRV_SDSD', 'HRV_CVNN', 'HRV_CVSD', 'HRV_MedianNN', 
+        'HRV_MadNN', 'HRV_MCVNN', 'HRV_IQRNN', 'HRV_SDRMSSD', 'HRV_Prc20NN', 'HRV_Prc80NN', 
+        'HRV_pNN50', 'HRV_pNN20', 'HRV_MinNN', 'HRV_MaxNN', 'HRV_HTI', 'HRV_TINN'],
+    FeatureDomain.FREQUENCY: [
+        'HRV_ULF', 'HRV_VLF', 'HRV_LF', 'HRV_HF', 'HRV_VHF', 'HRV_TP',
+       'HRV_LFHF', 'HRV_LFn', 'HRV_HFn', 'HRV_LnHF'
+       ],
+    FeatureDomain.NON_LINEAR: [
+        'HRV_SD1', 'HRV_SD2', 'HRV_SD1SD2', 'HRV_S', 'HRV_CSI', 'HRV_CVI',
+       'HRV_CSI_Modified', 'HRV_PIP', 'HRV_IALS', 'HRV_PSS', 'HRV_PAS',
+       'HRV_GI', 'HRV_SI', 'HRV_AI', 'HRV_PI', 'HRV_C1d', 'HRV_C1a',
+       'HRV_SD1d', 'HRV_SD1a', 'HRV_C2d', 'HRV_C2a', 'HRV_SD2d', 'HRV_SD2a',
+       'HRV_Cd', 'HRV_Ca', 'HRV_SDNNd', 'HRV_SDNNa', 'HRV_DFA_alpha1',
+       'HRV_MFDFA_alpha1_Width', 'HRV_MFDFA_alpha1_Peak',
+       'HRV_MFDFA_alpha1_Mean', 'HRV_MFDFA_alpha1_Max',
+       'HRV_MFDFA_alpha1_Delta', 'HRV_MFDFA_alpha1_Asymmetry',
+       'HRV_MFDFA_alpha1_Fluctuation', 'HRV_MFDFA_alpha1_Increment',
+       'HRV_DFA_alpha2', 'HRV_MFDFA_alpha2_Width', 'HRV_MFDFA_alpha2_Peak',
+       'HRV_MFDFA_alpha2_Mean', 'HRV_MFDFA_alpha2_Max',
+       'HRV_MFDFA_alpha2_Delta', 'HRV_MFDFA_alpha2_Asymmetry',
+       'HRV_MFDFA_alpha2_Fluctuation', 'HRV_MFDFA_alpha2_Increment',
+       'HRV_ApEn', 'HRV_SampEn', 'HRV_ShanEn', 'HRV_FuzzyEn', 'HRV_MSEn',
+       'HRV_CMSEn', 'HRV_RCMSEn', 'HRV_CD', 'HRV_HFD', 'HRV_KFD', 'HRV_LZC'
+       ]
+}
+
 
 class RPeak2HRV():
 
-    def get_hrv_features(self, input, windowing_method:str = None, time_header = "SystemTime", rri_header = "interbeat_interval", window_size = "60s", sampling_rate = 1000):
+    def get_hrv_features(self, input, windowing_method:str = None, time_header = "SystemTime", rri_header = "interbeat_interval", window_size = "60s", feature_domains = [FeatureDomain.TIME, FeatureDomain.FREQUENCY, FeatureDomain.NON_LINEAR],  sampling_rate = 1000):
         data = self._load_data(input)
         refined_data = self._refine_dataframe(input=data, sampling_rate=sampling_rate, time_header=time_header, rri_header=rri_header)
         if (windowing_method != None):
             windows = self._apply_windowing(data=refined_data, method=windowing_method, window_size= window_size)
-            # TODO: Implement hrv_time on windows
-            hrv_values = pd.DataFrame(columns=['window_start', 'window_end', 'HRV_MeanNN', 'HRV_SDNN', 'HRV_SDANN1', 'HRV_SDNNI1', 'HRV_SDANN2', 'HRV_SDNNI2', 'HRV_SDANN5', 'HRV_SDNNI5', 'HRV_RMSSD', 'HRV_SDSD', 'HRV_CVNN', 'HRV_CVSD', 'HRV_MedianNN', 'HRV_MadNN', 'HRV_MCVNN', 'HRV_IQRNN', 'HRV_SDRMSSD', 'HRV_Prc20NN', 'HRV_Prc80NN', 'HRV_pNN50', 'HRV_pNN20', 'HRV_MinNN', 'HRV_MaxNN', 'HRV_HTI', 'HRV_TINN'])
+            hrv_values = pd.DataFrame()
             for window in windows:
-                hrv_time = nk.hrv(self._convert_format(window), sampling_rate)
-                hrv_time['window_start'] = window.index[0]
-                hrv_time['window_end'] = window.index[-1]
-                hrv_values = pd.concat([hrv_values, hrv_time], ignore_index=True)
+                hrv_feature_values = self._calculate_features(window,feature_domains,sampling_rate)
+                hrv_feature_values['window_start'] = window.index[0]
+                hrv_feature_values['window_end'] = window.index[-1]
+                hrv_values = pd.concat([hrv_values, hrv_feature_values], ignore_index=True)
             return hrv_values
         else:
-            return nk.hrv_time(self._convert_format(refined_data), sampling_rate)
+            return self._calculate_features(refined_data, feature_domains, sampling_rate)
     
     def _load_data(self, input):
         if isinstance(input, str):
@@ -129,5 +164,18 @@ class RPeak2HRV():
             }
             return data_for_pipeline
 
+    def _calculate_features(self, data, feature_domains, sampling_rate):
+        # Calculate all possible features, then only keep needed
+        hrv_features = []
+        for feature in feature_domains:
+            try:
+                hrv_features.extend(FEATURE_COLS[feature])
+            except KeyError:
+                raise KeyError(f"{feature} is not a supported feature domain. feature_domains may only include 'time', 'freq' and 'non_lin'.")
+        result = nk.hrv(self._convert_format(data), sampling_rate)
+        return result[hrv_features]
+
+
+        
 
             

rpeaks_2_hrv_pipeline.py

@@ -1,5 +1,5 @@
 from transformers import Pipeline
-from .rpeaks2hrv import RPeak2HRV
+from .rpeaks2hrv import RPeak2HRV, FeatureDomain
 
 class RPeak2HRVPipeline(Pipeline):
     rpeak2HRVExtractor = RPeak2HRV()
@@ -16,10 +16,12 @@ class RPeak2HRVPipeline(Pipeline):
             preprocess_kwargs["rri_header"] = kwargs["rri_header"]
         if "window_size" in kwargs:
             preprocess_kwargs["window_size"] = kwargs["window_size"]
+        if "feature_domains" in kwargs:
+            preprocess_kwargs["feature_domains"] = kwargs["feature_domains"]
         return preprocess_kwargs, {}, {}
 
-    def preprocess(self, inputs, windowing_method:str = None, time_header = "SystemTime", rri_header = "interbeat_interval", window_size = "60s", sampling_rate = 1000):
-        return self.rpeak2HRVExtractor.get_hrv_features(inputs, windowing_method=windowing_method, time_header=time_header, rri_header=rri_header, window_size=window_size, sampling_rate=sampling_rate)
+    def preprocess(self, inputs, windowing_method:str = None, time_header = "SystemTime", rri_header = "interbeat_interval", window_size = "60s", feature_domains = [FeatureDomain.TIME, FeatureDomain.FREQUENCY, FeatureDomain.NON_LINEAR], sampling_rate = 1000):
+        return self.rpeak2HRVExtractor.get_hrv_features(inputs, windowing_method=windowing_method, time_header=time_header, rri_header=rri_header, window_size=window_size, feature_domains=feature_domains, sampling_rate=sampling_rate)
 
     def _forward(self, model_inputs):
         # currently empty as all preprocessing steps are performed by preprocess function