Fix testcase for logfilter and add option to extend the filter switch range by given number of pulses

eos/event_analysis.py

@@ -156,12 +156,14 @@ class FilterQzRange(EventDataAction):
 
 class FilterByLog(EventDataAction):
 
-    def __init__(self, filter_string, remove_switchpulse=False):
+    def __init__(self, filter_string, remove_switchpulse=False, pre_switch=0, post_switch=0):
         if filter_string.startswith('!'):
             filter_string = filter_string[1:]
             remove_switchpulse = True
         self.filter_string = filter_string
         self.remove_switchpulse = remove_switchpulse
+        self.pre_switch = pre_switch
+        self.post_switch = post_switch
 
     def perform_action(self, dataset: EventDatasetProtocol) -> None:
         filter_variable = None
@@ -184,8 +186,12 @@ class FilterByLog(EventDataAction):
             add_log_to_pulses(filter_variable, dataset)
         fltr_pulses = eval(self.filter_string, {filter_variable: dataset.data.pulses[filter_variable]})
         if self.remove_switchpulse:
-            switched = fltr_pulses[:-1] & ~fltr_pulses[1:]
-            fltr_pulses[:-1] &= ~switched
+            switched = np.logical_xor(fltr_pulses[:-1], fltr_pulses[1:])
+            fltr_pulses[:-1] &= np.logical_not(switched)
+            for i in range(self.pre_switch):
+                fltr_pulses[:-2-i] &= np.logical_not(switched[1+i:])
+            for i in range(self.post_switch):
+                fltr_pulses[1+i:-1] &= np.logical_not(switched[:-1-i])
         dataset.data.pulses.mask += EVENT_BITMASKS[filter_variable]*np.logical_not(fltr_pulses)
         goodTimeS = dataset.data.pulses.time[fltr_pulses]
         filter_e = np.logical_not(np.isin(dataset.data.events.wallTime, goodTimeS))

tests/test_event_handling.py

@@ -53,6 +53,7 @@ class MockEventData:
         pulses = np.recarray((packets.shape[0],), dtype=PULSE_TYPE)
         pulses.monitor = 1.0
         pulses.time = packets.time
+        pulses.mask = 0
 
         # proton current information with independent timing
         proton_current = np.recarray((50,), dtype=PC_TYPE)
@@ -552,27 +553,65 @@ class TestSimpleActions(TestCase):
 
         self._extract_walltime()
 
-        test_log = np.recarray(shape=(2,), dtype=np.dtype([('value', np.int32),
-                                                           ('time', np.int64)]))
-        test_log.time = [-5, self.d.data.pulses.time[100]+123]
-        test_log.value = [0, 1]
+        test_log = np.recarray(shape=(self.d.data.pulses.shape[0],),
+                               dtype=np.dtype([('value', np.int32), ('time', np.int64)]))
+        test_log.time = self.d.data.pulses.time
+        test_log.value = 0
+        test_log.value[10:-10] = 1
         self.d.data.device_logs['test_log'] = test_log
         action.perform_action(self.d)
-        self.assertEqual(self.d.data.pulses.shape[0], 101)
+        self.assertEqual(20, self.d.data.pulses.shape[0])
 
     def test_filter_by_log_switchpulse(self):
         action = FilterByLog("!test_log==0") | ApplyMask()
         self._extract_walltime()
 
-        test_log = np.recarray(shape=(2,), dtype=np.dtype([('value', np.int32),
-                                                           ('time', np.int64)]))
-        test_log.time = [-5, self.d.data.pulses.time[100]+123]
-        test_log.value = [0, 1]
+        test_log = np.recarray(shape=(self.d.data.pulses.shape[0],),
+                               dtype=np.dtype([('value', np.int32), ('time', np.int64)]))
+        test_log.time = self.d.data.pulses.time
+        test_log.value = 0
+        test_log.value[10:-10] = 1
         self.d.data.device_logs['test_log'] = test_log
         self.d.data.device_logs['check_log'] = test_log.copy()
         action.perform_action(self.d)
-        self.assertEqual(self.d.data.pulses.shape[0], 100)
+        self.assertEqual(19, self.d.data.pulses.shape[0])
         np.testing.assert_array_equal(
                 self.d.data.device_logs['test_log'],
                 self.d.data.device_logs['check_log'],
-                )
+                )
+
+    def test_filter_by_log_preswitch(self):
+        action = FilterByLog("!test_log==0", pre_switch=2) | ApplyMask()
+        self._extract_walltime()
+
+        test_log = np.recarray(shape=(self.d.data.pulses.shape[0],),
+                               dtype=np.dtype([('value', np.int32), ('time', np.int64)]))
+        test_log.time = self.d.data.pulses.time
+        test_log.value = 0
+        test_log.value[10:-10] = 1
+        self.d.data.device_logs['test_log'] = test_log
+        self.d.data.device_logs['check_log'] = test_log.copy()
+        action.perform_action(self.d)
+        self.assertEqual(17, self.d.data.pulses.shape[0])
+        np.testing.assert_array_equal(
+                self.d.data.device_logs['test_log'],
+                self.d.data.device_logs['check_log'],
+                )
+
+    def test_filter_by_log_postswitch(self):
+        action = FilterByLog("!test_log==0", post_switch=2) | ApplyMask()
+        self._extract_walltime()
+
+        test_log = np.recarray(shape=(self.d.data.pulses.shape[0],),
+                               dtype=np.dtype([('value', np.int32), ('time', np.int64)]))
+        test_log.time = self.d.data.pulses.time
+        test_log.value = 0
+        test_log.value[10:-10] = 1
+        self.d.data.device_logs['test_log'] = test_log
+        self.d.data.device_logs['check_log'] = test_log.copy()
+        action.perform_action(self.d)
+        self.assertEqual(17, self.d.data.pulses.shape[0])
+        np.testing.assert_array_equal(
+                self.d.data.device_logs['test_log'],
+                self.d.data.device_logs['check_log'],
+                )