Waveform cleaner

ctapipe/image/tests/test_waveform_cleaning.py

@@ -3,7 +3,8 @@
 
 from ctapipe.image.waveform_cleaning import (NullWaveformCleaner,
                                              CHECMWaveformCleanerAverage,
-                                             CHECMWaveformCleanerLocal)
+                                             CHECMWaveformCleanerLocal,
+                                             BaselineWaveformCleaner)
 
 
 def test_null_cleaner(example_event):
@@ -48,3 +49,24 @@ def test_checm_cleaner_local(example_event):
 
     assert_almost_equal(data_ped[0, 0, 0], -2.8, 1)
     assert_almost_equal(cleaned[0, 0, 0], -15.9, 1)
+
+
+def test_baseline_cleaner():
+
+    # waveform : first 20 samples = 0, second 20 samples = 10
+    waveform = np.full((2, 1855, 40), 10)
+    waveform[:, :, 0:20] = 0
+
+    cleaner = BaselineWaveformCleaner()
+
+    cleaner.baseline_start = 0
+    cleaner.baseline_end = 20
+    cleaned = cleaner.apply(waveform)
+    assert (cleaned.mean() == 5)
+
+    cleaner.baseline_start = 20
+    cleaner.baseline_end = 40
+    cleaned = cleaner.apply(waveform)
+    assert (cleaned.mean() == -5)
+
+

ctapipe/image/waveform_cleaning.py

@@ -13,7 +13,7 @@
                                              LocalPeakIntegrator)
 
 __all__ = ['WaveformCleanerFactory', 'CHECMWaveformCleanerAverage',
-           'CHECMWaveformCleanerLocal',
+           'CHECMWaveformCleanerLocal', 'BaselineWaveformCleaner',
            'NullWaveformCleaner']
 
 
@@ -67,6 +67,24 @@ def apply(self, waveforms):
         return waveforms
 
 
+class BaselineWaveformCleaner(WaveformCleaner):
+    """
+    Basic waveform cleaner that subtracts the waveform baseline
+    estimated as the mean waveform value in the interval [baseline_start,baseline_end]
+    """
+    baseline_start = Int(0, help='Start sample for baseline estimation').tag(config=True)
+
+    baseline_end = Int(10, help='End sample for baseline estimation').tag(config=True)
+
+    def apply(self, waveforms):
+        # Subtract baseline
+        baseline_corrected = waveforms - np.mean(
+            waveforms[:, :, self.baseline_start:self.baseline_end], axis=2
+        )[:, :, None]
+
+        return baseline_corrected
+
+
 class CHECMWaveformCleaner(WaveformCleaner):
     """
     Waveform cleaner used by CHEC-M.