Merge pull request #414 from Mentalab-hub/imp-calibration-32ch-pro

Impedance calibration update for pro devices

Author: salman2135

src/explorepy/packet.py

@@ -43,6 +43,7 @@ class PACKET_ID(IntEnum):
     PUSHMARKER = 194
     CALIBINFO = 195
     CALIBINFO_USBC = 197
+    CALIBINFO_PRO = 196
     TRIGGER_OUT = 177  # Trigger-out of Explore device
     TRIGGER_IN = 178  # Trigger-in to Explore device
     VERSION_INFO = 199
@@ -225,15 +226,19 @@ def int32_to_status(data):
         }
         return status
 
-    def calculate_impedance(self, imp_calib_info):
+    def calculate_impedance(self, imp_calib_info, index=None):
         """calculate impedance with the help of impedance calibration info
 
         Args:
             imp_calib_info (dict): dictionary of impedance calibration info including slope, offset and noise level
 
         """
-        scale = imp_calib_info["slope"]
-        offset = imp_calib_info["offset"]
+        if index is None:
+            scale = imp_calib_info["slope"]
+            offset = imp_calib_info["offset"]
+        else:
+            scale = imp_calib_info["slope"][index]
+            offset = imp_calib_info["offset"][index]
         self.imp_data = np.round(
             (self.get_ptp()
              - imp_calib_info["noise_level"]) * scale / 1.0e6 - offset,
@@ -702,35 +707,57 @@ def __str__(self):
 
 
 class CalibrationInfoBase(Packet):
-    @abc.abstractmethod
-    def _convert(self, bin_data, offset_multiplier=0.001):
-        slope = np.frombuffer(bin_data,
-                              dtype=np.dtype(np.uint16).newbyteorder("<"),
-                              count=1,
-                              offset=0).item()
-        self.slope = slope * 10.0
-        offset = np.frombuffer(bin_data,
-                               dtype=np.dtype(np.uint16).newbyteorder("<"),
-                               count=1,
-                               offset=2).item()
-        self.offset = offset * offset_multiplier
+    """Base class for calibration packets"""
+
+    channels = 1
+    offset_multiplier = 0.001
+
+    def __init__(self, timestamp, payload, time_offset=0):
+        # Must exist before Packet.__init__ calls _convert()
+        self.slope = []
+        self.offset = []
+        super().__init__(timestamp, payload, time_offset)
+
+    def _convert(self, bin_data):
+        dtype_u16 = np.dtype("<u2")
+
+        for i in range(self.channels):
+            base = i * 4
+
+            slope = np.frombuffer(
+                bin_data,
+                dtype=dtype_u16,
+                count=1,
+                offset=base
+            ).item()
+            self.slope.append(slope * 10.0)
+
+            offset = np.frombuffer(
+                bin_data,
+                dtype=dtype_u16,
+                count=1,
+                offset=base + 2
+            ).item()
+            self.offset.append(offset * self.offset_multiplier)
 
     def get_info(self):
-        """Get calibration info"""
         return {"slope": self.slope, "offset": self.offset}
 
     def __str__(self):
-        return "calibration info: slope = " + str(self.slope) + "\toffset = " + str(self.offset)
+        return f"calibration info: slope = {self.slope}\toffset = {self.offset}"
 
 
 class CalibrationInfo(CalibrationInfoBase):
-    def _convert(self, bin_data):
-        super()._convert(bin_data, offset_multiplier=0.001)
+    offset_multiplier = 0.001
 
 
 class CalibrationInfo_USBC(CalibrationInfoBase):
-    def _convert(self, bin_data):
-        super()._convert(bin_data, offset_multiplier=0.01)
+    offset_multiplier = 0.01
+
+
+class CalibrationInfoPro(CalibrationInfoBase):
+    channels = 4
+    offset_multiplier = 0.01
 
 
 class BleImpedancePacket(EEG98_USBC):
@@ -751,6 +778,19 @@ def populate_packet_with_data(self, ble_packet_list):
                 data_array = np.concatenate((data_array, data), axis=1)
         self.data = data_array
 
+    def resize_packet(self, full_data, index):
+        self.data = full_data[index * 8: index * 8 + 8, :]
+
+    def populate_data_1d(self, ble_packet_list):
+        data_array = None
+        for i in range(len(ble_packet_list)):
+            _, data = ble_packet_list[i].get_data()
+            if data_array is None:
+                data_array = data
+            else:
+                data_array = np.concatenate((data_array, data), axis=0)
+        self.data = data_array
+
 
 class VersionInfoPacket(Packet):
     def __init__(self, timestamp, payload, time_offset=0):
@@ -792,6 +832,7 @@ def __str__(self):
     PACKET_ID.CMDSTAT: CommandStatus,
     PACKET_ID.CALIBINFO: CalibrationInfo,
     PACKET_ID.CALIBINFO_USBC: CalibrationInfo_USBC,
+    PACKET_ID.CALIBINFO_PRO: CalibrationInfoPro,
     PACKET_ID.PUSHMARKER: PushButtonMarker,
     PACKET_ID.TRIGGER_IN: TriggerIn,
     PACKET_ID.TRIGGER_OUT: TriggerOut,

src/explorepy/stream_processor.py

@@ -23,8 +23,7 @@
 from explorepy.filters import ExGFilter
 from explorepy.packet import (
     EEG,
-    CalibrationInfo,
-    CalibrationInfo_USBC,
+    CalibrationInfoBase,
     CommandRCV,
     CommandStatus,
     DeviceInfo,
@@ -173,8 +172,8 @@ def group_packets_by_base_type(self, packet_batch: List[Tuple]):
                 base_class = Environment
             elif isinstance(packet, EventMarker):
                 base_class = EventMarker
-            elif isinstance(packet, CalibrationInfo):
-                base_class = CalibrationInfo
+            elif isinstance(packet, CalibrationInfoBase):
+                base_class = CalibrationInfoBase
             elif isinstance(packet, PacketBIN):
                 base_class = PacketBIN
             else:
@@ -225,8 +224,8 @@ def process_batch(self, packet_batch: List[Tuple]):
             self._process_marker_batch(sorted_eeg_packets)
 
         # Process calibration info packets
-        if CalibrationInfo in grouped_packets:
-            self._process_calib_info_batch(grouped_packets[CalibrationInfo])
+        if CalibrationInfoBase in grouped_packets:
+            self._process_calib_info_batch(grouped_packets[CalibrationInfoBase])
 
         # Process binary packets
         if PacketBIN in grouped_packets:
@@ -348,7 +347,7 @@ def process(self, packet):
             self.dispatch(topic=TOPICS.env, packet=packet)
         elif isinstance(packet, EventMarker):
             self.dispatch(topic=TOPICS.marker, packet=packet)
-        elif isinstance(packet, CalibrationInfo) or isinstance(packet, CalibrationInfo_USBC):
+        elif isinstance(packet, CalibrationInfoBase):
             self.imp_calib_info = packet.get_info()
         elif not packet:
             self.is_connected = False
@@ -449,11 +448,12 @@ def configure_device(self, cmd):
         self.start_cmd_process_thread()
         return self._device_configurator.change_setting(cmd)
 
-    def imp_initialize(self, notch_freq):
+    def imp_initialize(self, notch_freq, calibration=False):
         """Activate impedance mode in the device"""
         logger.info("Starting impedance measurement mode...")
         cmd = ZMeasurementEnable()
         if self.configure_device(cmd):
+            self.imp_calib_info['calibration'] = calibration
             self.imp_calculator = ImpedanceMeasurement(device_info=self.device_info,
                                                        calib_param=self.imp_calib_info,
                                                        notch_freq=notch_freq)

src/explorepy/tools.py

@@ -438,8 +438,7 @@ def __init__(self, device_info, stream_name=None):
         self.marker_outlet = None
         self.exg_outlet = None
         self.orn_outlet = None
-        self.adc_mask = SettingsManager(
-            device_info["device_name"]).get_adc_mask()
+        self.adc_mask = SettingsManager(device_info["device_name"]).get_adc_mask()
         if len(SettingsManager(device_info["device_name"]).get_channel_names()) == len(self.adc_mask):
             self.channel_names = SettingsManager(device_info["device_name"]).get_channel_names()
         else:
@@ -544,64 +543,92 @@ def __init__(self, device_info, calib_param, notch_freq):
         self.packet_buffer = []
 
     def _add_filters(self):
-        bp_freq = self._device_info['sampling_rate'] / 4 - \
-            1.5, self._device_info['sampling_rate'] / 4 + 1.5
-        noise_freq = self._device_info['sampling_rate'] / \
-            4 + 2.5, self._device_info['sampling_rate'] / 4 + 5.5
+        s_rate = 250
+        center = s_rate / 4
+
+        bp_freq = (center - 1.5, center + 1.5)
+        noise_freq = (center + 2.5, center + 5.5)
+
         settings_manager = SettingsManager(self._device_info["device_name"])
-        settings_manager.load_current_settings()
-        n_chan = settings_manager.settings_dict[settings_manager.channel_count_key]
-        # Temporary fix for 16/32 channel filters
-        if not is_explore_pro_device():
-            if n_chan >= 16:
-                n_chan = 32
-        self._filters['notch'] = ExGFilter(cutoff_freq=self._notch_freq,
-                                           filter_type='notch_imp',
-                                           s_rate=self._device_info['sampling_rate'],
-                                           n_chan=n_chan)
-
-        self._filters['demodulation'] = ExGFilter(cutoff_freq=bp_freq,
-                                                  filter_type='bandpass',
-                                                  s_rate=self._device_info['sampling_rate'],
-                                                  n_chan=n_chan)
-
-        self._filters['base_noise'] = ExGFilter(cutoff_freq=noise_freq,
-                                                filter_type='bandpass',
-                                                s_rate=self._device_info['sampling_rate'],
-                                                n_chan=n_chan)
+        n_chan = settings_manager.get_channel_count()
+        self.adc_count = n_chan // 8
+
+        def make_filter(filter_type, cutoff, channels):
+            return ExGFilter(
+                cutoff_freq=cutoff,
+                filter_type=filter_type,
+                s_rate=s_rate,
+                n_chan=channels,
+            )
+
+        if not self._calib_param['calibration']:
+
+            self._filters['notch'] = [
+                make_filter('notch_imp', self._notch_freq, 8)
+                for _ in range(self.adc_count)
+            ]
+            self._filters['demodulation'] = [
+                make_filter('bandpass', bp_freq, 8)
+                for _ in range(self.adc_count)
+            ]
+            self._filters['base_noise'] = [
+                make_filter('bandpass', noise_freq, 8)
+                for _ in range(self.adc_count)
+            ]
+        else:
+            self._filters['notch'] = make_filter('notch_imp', self._notch_freq, n_chan)
+            self._filters['demodulation'] = make_filter('bandpass', bp_freq, n_chan)
+            self._filters['base_noise'] = make_filter('bandpass', noise_freq, n_chan)
 
     def measure_imp(self, packet):
         """Compute electrode impedances
         """
+        if self._calib_param['calibration']:
+            return None
         self.packet_buffer.append(packet)
 
         if len(self.packet_buffer) < 16:
             return None
         else:
-            timestamp, _ = self.packet_buffer[0].get_data()
+            timestamp, data = self.packet_buffer[0].get_data()
             resized_packet = BleImpedancePacket(
                 timestamp=timestamp, payload=None)
             resized_packet.populate_packet_with_data(self.packet_buffer)
             self.packet_buffer.clear()
-            temp_packet = self._filters['notch'].apply(
-                input_data=resized_packet, in_place=False)
-            self._calib_param['noise_level'] = self._filters['base_noise']. \
-                apply(input_data=temp_packet, in_place=False).get_ptp()
-            self._filters['demodulation'].apply(
-                input_data=temp_packet, in_place=True
-            ).calculate_impedance(self._calib_param)
-            return temp_packet
+            imp_packet_buffer = []
+            if isinstance(self._filters['notch'], list):
+                temp_packet = None
+                # slice packet and feed it to the filters
+                for i in range(self.adc_count):
+                    sliced_packet = BleImpedancePacket(
+                        timestamp=timestamp, payload=None)
+                    sliced_packet.resize_packet(resized_packet.get_data()[1], i)
+                    temp_packet = self._filters['notch'][i].apply(
+                        input_data=sliced_packet, in_place=False)
+                    self._calib_param['noise_level'] = self._filters['base_noise'][i]. \
+                        apply(input_data=temp_packet, in_place=False).get_ptp()
+                    self._filters['demodulation'][i].apply(
+                        input_data=temp_packet, in_place=True
+                    )
+                    temp_packet.calculate_impedance(self._calib_param, index=i)
+                    temp_packet.data = temp_packet.imp_data
+                    imp_packet_buffer.append(temp_packet)
+                imp_packet = BleImpedancePacket(timestamp=timestamp, payload=None)
+                imp_packet.populate_data_1d(imp_packet_buffer)
+                imp_packet.imp_data = imp_packet.data
+
+                return imp_packet
 
 
 def find_free_port():
-    """Find a free port on the localhost
+    """Find a free port on the localhost)
+        free_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
 
     Returns:
         int: Port number
     """
     with closing(socket.socket(socket.AF_INET, socket.SOCK_STREAM)) as free_socket:
         free_socket.bind(('localhost', 0))
-        free_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)
         port_number = free_socket.getsockname()[1]
         return port_number