analyzers: define sos filter

Signed-off-by: Timofey Titovets <nefelim4ag@gmail.com>

scripts/motan/analyzers.py

@@ -197,6 +197,77 @@ class GenSmoothed:
         return data
 AHandlers["smooth"] = GenSmoothed
 
+class GenSOSFilter:
+    ParametersMin = 5
+    ParametersMax = 6
+    DataSets = [
+        ('sos(<dataset>,{filt,filtfilt},highpass,<order>,<hz>)',
+         'SOS highpass filtered dataset'),
+        ('sos(<dataset>,{filt,filtfilt},lowpass,<order>,<hz>)',
+         'SOS lowpass filtered dataset'),
+        ('sos(<dataset>,{filt,filtfilt},bandpass,<order>,<lowhz>,<highhz>)',
+         'SOS lowpass filtered dataset'),
+        ('sos(<dataset>,{filt,filtfilt},notch,<hz>,<quality>)',
+         'SOS notch filtered dataset'),
+    ]
+    def __init__(self, amanager, name_parts):
+        self.amanager = amanager
+        self.source = name_parts[1]
+
+        from numpy import array
+        from scipy.signal import sosfiltfilt, sosfilt, sosfilt_zi
+        self.array = array
+        self.type = name_parts[2]
+        if self.type == "filt":
+            self.sosfilt = sosfilt
+            self.sosfilt_zi = sosfilt_zi
+        elif self.type == "filtfilt":
+            self.sosfilt = sosfiltfilt
+        else:
+            amanager.error("Unknown sos type '%s'" % (self.type,))
+
+        filter_type = name_parts[3]
+        amanager.setup_dataset(self.source)
+        inv_seg_time = 1.0 / amanager.get_segment_time()
+
+        if filter_type in ('highpass', 'lowpass', 'bandpass'):
+            from scipy.signal import butter
+            order = int(name_parts[4])
+            cutoff = float(name_parts[5])
+            desc = "%.fHz" % (cutoff)
+            if filter_type == 'bandpass':
+                cutoff = (cutoff, float(name_parts[6]))
+                desc = "%.1f..%.1fHz" % cutoff
+            self.sos = butter(order, cutoff, filter_type, output='sos',
+                              fs=inv_seg_time)
+            self.filter_desc = "%s %s order %d" % (filter_type, desc, order)
+        elif filter_type == 'notch':
+            from scipy.signal import iirnotch, tf2sos
+            freq = float(name_parts[4])
+            quality = float(name_parts[5])
+            b, a = iirnotch(freq, Q=quality, fs=inv_seg_time)
+            self.sos = tf2sos(b, a)
+            self.filter_desc = "notch %.1fHz Q: %.1f" % (freq, quality)
+        else:
+            raise amanager.error("Unknown filter type '%s'" % (filter_type,))
+
+    def get_label(self):
+        label = self.amanager.get_label(self.source)
+        lname = "SOS %s (%s)" % (self.filter_desc, label['label'])
+        return {'label': lname, 'units': label['units']}
+
+    def generate_data(self):
+        data = self.amanager.get_datasets()[self.source]
+        data_array = self.array(data)
+        if self.type == "filt":
+            zi = self.sosfilt_zi(self.sos) * data_array[0]
+            filtered, _ = self.sosfilt(self.sos, data_array, zi=zi)
+        else:
+            filtered = self.sosfilt(self.sos, data_array)
+        return filtered.tolist()
+
+AHandlers["sos"] = GenSOSFilter
+
 # Calculate a kinematic stepper position from the toolhead requested position
 class GenKinematicPosition:
     ParametersMin = ParametersMax = 1