From c72e1df205f02c3298dfd8cf83bb5d4f7388f9e9 Mon Sep 17 00:00:00 2001
From: Olivier Winter <olivier.winter@hotmail.fr>
Date: Tue, 8 Nov 2022 15:01:41 +0000
Subject: [PATCH] spacer detection method

---
 iblrig/spacer.py            | 20 ++++++++++++++++++++
 test_iblrig/test_spacers.py | 15 +++++++++++++++
 2 files changed, 35 insertions(+)

diff --git a/iblrig/spacer.py b/iblrig/spacer.py
index 1c6301b59..e39283c42 100644
--- a/iblrig/spacer.py
+++ b/iblrig/spacer.py
@@ -95,3 +95,23 @@ def add_spacer_states(self, sma=None, next_state="exit"):
                 state_change_conditions={"Tup": next_loop},
                 output_actions=[],
             )
+
+    def find_spacers(self, signal, threshold=0.9, fs=1000):
+        """
+        Find spacers in a voltage time serie. Assumes that the signal is a digital signal between 0 and 1
+        :param signal:
+        :param threshold:
+        :param fs:
+        :return:
+        """
+        template = self.generate_template(fs=fs)
+        xcor = np.correlate(signal, template, mode="full") / np.sum(template)
+        idetect = np.where(xcor > threshold)[0]
+        iidetect = np.cumsum(np.diff(idetect, prepend=0) > 1)
+        nspacers = iidetect[-1]
+        tspacer = np.zeros(nspacers)
+        for i in range(nspacers):
+            ispacer = idetect[iidetect == i + 1]
+            imax = np.argmax(xcor[ispacer])
+            tspacer[i] = (ispacer[imax] - template.size + 1) / fs
+        return tspacer
diff --git a/test_iblrig/test_spacers.py b/test_iblrig/test_spacers.py
index 701da9111..87d217c78 100644
--- a/test_iblrig/test_spacers.py
+++ b/test_iblrig/test_spacers.py
@@ -17,3 +17,18 @@ def test_spacer(self):
         ac[sig.size-100: sig.size + 100] = 0  # remove the main peak
         # the autocorrelation side lobes should be less than 30%
         assert np.max(ac) < .3
+
+    def test_find_spacers(self):
+        """
+        Generates a fake signal with 2 spacers and finds them
+        :return:
+        """
+        fs = 1000
+        spacer = Spacer(dt_start=.02, dt_end=.4, n_pulses=8, tup=.05)
+        start_times = [4.38, 96.58]
+        template = spacer.generate_template(fs)
+        signal = np.zeros(int(start_times[-1] * fs + template.size * 2))
+        for start_time in start_times:
+            signal[int(start_time * fs): int(start_time * fs) + template.size] = template
+        spacer_times = spacer.find_spacers(signal, fs=fs)
+        np.testing.assert_allclose(spacer_times, start_times)