improvements to the realtime module

teonbrooks · teonbrooks · commit 6763f7fdfeba · 2019-04-11T20:57:13.000-04:00
currently the test is breaking when it comes to using the RtEpochs object.
diff --git a/mne/realtime/base_client.py b/mne/realtime/base_client.py
@@ -133,6 +133,12 @@ def register_receive_callback(self, callback):
         if callback not in self._recv_callbacks:
             self._recv_callbacks.append(callback)
 
+    def start(self):
+        """Start the client."""
+        self.__enter__()
+
+        return self
+
     def start_receive_thread(self, nchan):
         """Start the receive thread.
 
@@ -150,6 +156,12 @@ def start_receive_thread(self, nchan):
             self._recv_thread.daemon = True
             self._recv_thread.start()
 
+    def stop(self):
+        """Stop the client."""
+        self.__exit__()
+
+        return self
+
     def stop_receive_thread(self, stop_measurement=False):
         """Stop the receive thread.
 
diff --git a/mne/realtime/lsl_client.py b/mne/realtime/lsl_client.py
@@ -74,11 +74,8 @@ def get_data_as_epoch(self, n_samples=1024, picks=None):
 
     def iter_raw_buffers(self):
         """Return an iterator over raw buffers."""
-        pylsl = _check_pylsl_installed(strict=True)
-        inlet = pylsl.StreamInlet(self.client)
-
         while True:
-            samples, _ = inlet.pull_chunk(max_samples=self.buffer_size)
+            samples, _ = self.client.pull_chunk(max_samples=self.buffer_size)
 
             yield np.vstack(samples).T
 
diff --git a/mne/realtime/mock_lsl_stream.py b/mne/realtime/mock_lsl_stream.py
@@ -31,63 +31,68 @@ class MockLSLStream:
 
     def __init__(self, host, n_channels=8, ch_type="eeg", sfreq=100,
                  testing=False):
-        self.host = host
-        self.n_channels = n_channels
-        self.ch_type = ch_type
-        self.sfreq = sfreq
-        self.streaming = False
-        self.testing = testing
+        self._host = host
+        self._n_channels = n_channels
+        self._ch_type = ch_type
+        self._sfreq = sfreq
+        self._testing = testing
+        self._streaming = False
 
     def start(self):
         """Start a mock LSL stream."""
-        pylsl = _check_pylsl_installed(strict=True)
-        self.streaming = True
-        info = pylsl.StreamInfo('MNE', self.ch_type.upper(), self.n_channels,
-                                self.sfreq, 'float32', self.host)
-        info.desc().append_child_value("manufacturer", "MNE")
-        channels = info.desc().append_child("channels")
-        for c_id in range(1, self.n_channels + 1):
-            channels.append_child("channel") \
-                    .append_child_value("label", "MNE {:03d}".format(c_id)) \
-                    .append_child_value("type", self.ch_type.lower()) \
-                    .append_child_value("unit", "microvolt")
-
-        # next make an outlet
-        outlet = pylsl.StreamOutlet(info)
-
         print("now sending data...")
-        self.process = Process(target=self._initiate_stream, args=(outlet,))
+        self.process = Process(target=self._initiate_stream)
         self.process.daemon = True
         self.process.start()
 
         return self
 
-    def close(self):
+    def stop(self):
         """Stop a mock LSL stream."""
+        self._streaming = False
         self.process.terminate()
 
         print("Stopping stream...")
 
         return self
 
-    def _initiate_stream(self, outlet):
+    def _initiate_stream(self):
+        # outlet needs to be made on the same process
+        pylsl = _check_pylsl_installed(strict=True)
+        self._streaming = True
+        info = pylsl.StreamInfo(name='MNE', type=self._ch_type.upper(),
+                                channel_count=self._n_channels,
+                                nominal_srate=self._sfreq,
+                                channel_format='float32', source_id=self._host)
+        info.desc().append_child_value("manufacturer", "MNE")
+        channels = info.desc().append_child("channels")
+        for c_id in range(1, self._n_channels + 1):
+            channels.append_child("channel") \
+                    .append_child_value("label", "MNE {:03d}".format(c_id)) \
+                    .append_child_value("type", self._ch_type.lower()) \
+                    .append_child_value("unit", "microvolt")
+
+        # next make an outlet
+        outlet = pylsl.StreamOutlet(info)
+
+        # let's make some data
         counter = 0
         trigger = 0
-        while True:
-            sample = counter % self.sfreq
+        while self._streaming:
+            sample = counter % self._sfreq
             # let's bound trigger to be between 1 and 10 so the max cycle
             # is ten seconds
-            if trigger == 10:
-                trigger = 0
+            trigger = 0 if trigger == 10 else trigger
             if sample == 0:
                 trigger += 1
 
-            if not self.testing:
-                const = np.sin(2 * np.pi * sample / self.sfreq) * 1e-6
-                mysample = rand(self.n_channels).dot(const).tolist()
+            if self._testing:
+                const = trigger
             else:
-                mysample = np.ones(self.n_channels).dot(trigger).tolist()
+                const = np.sin(2 * np.pi * sample / self._sfreq) * 1e-6
+
+            mysample = rand(self._n_channels).dot(const).tolist()
             # now send it and wait for a bit
             outlet.push_sample(mysample)
             counter += 1
-            time.sleep(self.sfreq**-1)
+            time.sleep(self._sfreq**-1)
diff --git a/mne/realtime/tests/test_lsl_client.py b/mne/realtime/tests/test_lsl_client.py
@@ -2,14 +2,26 @@
 #
 # License: BSD (3-clause)
 from random import random as rand
+import numpy as np
 
 from mne.realtime import LSLClient, MockLSLStream, RtEpochs
 from mne.utils import run_tests_if_main, requires_pylsl
 from mne import create_info
+from mne.io import constants
 
 
 host = 'myuid34234'
-
+event_id = 5
+n_channels = 8
+sfreq = 100
+tmin = -0.1
+tmax = 0.5
+ch_types = ["eeg" for n in range(n_channels - 1)] + ['stim']
+ch_names = ["MNE {:03d}".format(ch_id) for ch_id
+            in range(1, n_channels + 1)]
+stim_channel = ch_names[-1]
+
+info = create_info(ch_names, sfreq, ch_types)
 
 @requires_pylsl
 def test_lsl_client():
@@ -21,47 +33,54 @@ def test_lsl_client():
     stream = MockLSLStream(host, n_channels, testing=True)
     stream.start()
 
-    with LSLClient(info=None, host=host, wait_max=wait_max) as client:
+    with LSLClient(info=info, host=host, wait_max=wait_max) as client:
         client_info = client.get_measurement_info()
-
-        assert ([ch["ch_name"] for ch in client_info["chs"]] ==
-                ["MNE {:03d}".format(ch_id) for ch_id in
-                 range(1, n_channels + 1)])
-
         epoch = client.get_data_as_epoch(n_samples=n_samples)
-        assert n_channels, n_samples == epoch.get_data().shape[1:]
 
-    stream.close()
+    assert client_info['nchan'] == n_channels
+    assert ([ch["ch_name"] for ch in client_info["chs"]] ==
+            [ch_name for ch_name in ch_names])
+    assert any([constants.FIFF.FIFFV_STIM_CH == ch['kind']
+                for ch in info['chs']])
+    assert n_channels, n_samples == epoch.get_data().shape[1:]
+
+    stream.stop()
 
 
 def test_lsl_rt_epochs():
     """Test the functionality of the LSL Client with RtEpochs."""
-    event_id = 5
-    n_channels = 8
-    sfreq = 100
-    tmin = -0.1
-    tmax = 0.5
-    ch_types = ["eeg" for n in range(n_channels - 1)] + ['stim']
-    ch_names = ["MNE {:03d}".format(ch_id) for ch_id
-                in range(1, n_channels + 1)]
-    stim_channel = ch_names[-1]
-
     stream = MockLSLStream(host, n_channels=n_channels, ch_type="eeg",
                            sfreq=sfreq, testing=True)
     stream.start()
 
-    info = create_info(ch_names, sfreq, ch_types)
-    with LSLClient(info=info, host=host) as client:
-        epochs_rt = RtEpochs(client, event_id, tmin, tmax, stim_channel)
-        epochs_rt.start()
-        time.sleep(10)
-        epochs.stop(stop_receive_thread=True)
+    try:
+        data = None
+        events_ids = None
+        with LSLClient(info=info, host=host) as client:
+            epochs_rt = RtEpochs(client, event_id, tmin, tmax, stim_channel)
+            epochs_rt.start()
+            time.sleep(15)
+
+            for ev_num, ev in enumerate(epochs_rt.iter_evoked()):
+                if ev_num == 0:
+                    data = ev.data[None, :, :]
+                    events_ids = int(
+                        ev.comment)  # comment attribute contains event_id
+                else:
+                    data = np.concatenate(
+                        (data, ev.data[None, :, :]), axis=0)
+                    events_ids = np.append(events_ids,
+                                           int(ev.comment))
+
+            epochs_rt.stop(stop_receive_thread=True)
 
         data = epochs_rt.get_data()
+
         n_samples_prestim = np.abs(tmin * sfreq)
         n_sample_poststim = tmax * sfreq
 
-        data_channels, data_samples = data.get_data().shape[1:]
+        data_channels, data_samples = data.shape[1:]
+
         assert data_channels == n_channels
         assert data_samples == n_samples_prestim + n_sample_poststim
 
@@ -72,4 +91,7 @@ def test_lsl_rt_epochs():
 
         assert data == data_expected
 
+    finally:
+        stream.stop()
+
 run_tests_if_main()
