Dev: to edf

mcc-flow/merge_records.py

+from dataclasses import dataclass, field
 from datetime import datetime
 from pathlib import Path
 from struct import unpack
 
 import numpy as np
 from matplotlib import pyplot as plt
+from pyedflib import EdfWriter, FILETYPE_EDFPLUS
 
+HEADER_FS = 1
+EEG_EMG_FS = 250
+GYRO_ACC_FS = 50
+SOUND_FS = 6000
+GSR_FS = 10
 
+EEG_EMG_SCALE = 0.04470348358154  # (mV)
+GYRO_SCALE = 0.0152587890625  # (deg/sec)
+ACC_SCALE = 0.00006103515625  # (g)
+
+
+@dataclass
+class PacketData:
+    data: np.ndarray
+    type: str
+    unit: str
+    fs: int
+
+    @property
+    def chs(self):
+        return self.data.shape[0]
+
+    def __add__(self, other):
+        assert self.type == other.type, 'Cannot combine PacketDatas with different signal types'
+        assert self.unit == other.unit, 'Cannot combine PacketDatas with different units'
+        assert self.fs == other.fs, 'Cannot combine PacketDatas with different sampling frequency'
+        assert self.chs == other.chs, 'Cannot combine PacketDatas with different number of channels'
+        self.data = np.concatenate((self.data, other.data), axis=-1)
+        return self
+
+
+@dataclass
 class FlowPacket:
-    header_fs = 1
-    signal_fs = 250
-    gyro_acc_fs = 50
-    sound_fs = 6000
-    grs_fs = 10
-
-    def __init__(self, date_, flow, hr, spo2,
-                 # vbat_eeg, vbat_arm,
-                 eeg, eeg_giro, eeg_acc,
-                 emg, emg_giro, emg_acc,
-                 sound, gsr):
-        self.date = date_
-        self.duration = 1
-        self.flow = [flow]
-        self.hr = [hr]
-        self.spo2 = [spo2]
-
-        self.eeg = eeg
-        self.eeg_giro = eeg_giro
-        self.eeg_acc = eeg_acc
-        self.emg = emg
-        self.emg_giro = emg_giro
-        self.emg_acc = emg_acc
-        self.sound = sound
-        self.gsr = gsr
+    date: datetime
+    duration: int = field(init=False, default=1)
+    flow: PacketData
+    hr: PacketData
+    spo2: PacketData
+    eeg: PacketData
+    eeg_giro: PacketData
+    eeg_acc: PacketData
+    emg: PacketData
+    emg_giro: PacketData
+    emg_acc: PacketData
+    sound: PacketData
+    gsr: PacketData
 
     def __add__(self, other):
-        self.flow.extend(other.flow)
-        self.hr.extend(other.hr)
-        self.spo2.extend(other.spo2)
-        self.eeg = np.concatenate((self.eeg, other.eeg), axis=-1)
-        self.eeg_giro = np.concatenate((self.eeg_giro, other.eeg_giro), axis=-1)
-        self.eeg_acc = np.concatenate((self.eeg_acc, other.eeg_acc), axis=-1)
-        self.emg = np.concatenate((self.emg, other.emg), axis=-1)
-        self.emg_giro = np.concatenate((self.emg_giro, other.emg_giro), axis=-1)
-        self.emg_acc = np.concatenate((self.emg_acc, other.emg_acc), axis=-1)
-        self.sound = np.concatenate((self.sound, other.sound), axis=-1)
-        self.gsr = np.concatenate((self.gsr, other.gsr), axis=-1)
+        self.flow += other.flow
+        self.hr += other.hr
+        self.spo2 += other.spo2
+        self.eeg += other.eeg
+        self.eeg_giro += other.eeg_giro
+        self.eeg_acc += other.eeg_acc
+        self.emg += other.emg
+        self.emg_giro += other.emg_giro
+        self.emg_acc += other.emg_acc
+        self.sound += other.sound
+        self.gsr += other.gsr
         self.duration += other.duration
         return self
 
-    def _plot(self, data, title):
+    def _plot(self, pkt_data, title):
+        data = pkt_data.data
         plt.figure()
         plt.suptitle(title)
         chs = data.shape[0]
         for i in range(chs):
             plt.subplot(chs, 1, i + 1)
             plt.plot(np.linspace(0, self.duration, data.shape[1]), data[i, :])
-            plt.title(f'ch{i + 1}')
+            if chs > 1:
+                plt.title(f'ch{i + 1}')
         plt.xlabel('time (s)')
 
     def plot(self):
         self._plot(self.eeg, 'EEG')
         self._plot(self.eeg_giro, 'EEG Gyro')
         self._plot(self.emg, 'EMG')
-
+        self._plot(self.spo2, 'SPO2')
         plt.show()
 
+    def to_edf(self, filename):
+        to_save = [self.eeg, self.eeg_giro, self.eeg_acc,
+                   self.emg, self.emg_giro, self.emg_acc,
+                   self.flow, self.hr, self.spo2,
+                   self.sound, self.gsr]
+
+        with EdfWriter(filename, sum(pkt.chs for pkt in to_save), FILETYPE_EDFPLUS) as edf:
+            data_buffer = []
+            sig = 0
+            for pkt in to_save:
+                for i, ch_dat in enumerate(pkt.data):
+                    print(sig)
+                    # edf.setPhysicalMaximum(sig, 2**7 * EEG_EMG_SCALE)
+                    # edf.setPhysicalMinimum(sig, -2**7 * EEG_EMG_SCALE)
+                    # edf.setDigitalMaximum(sig, 2**7)
+                    # edf.setDigitalMinimum(sig, -2**7)
+                    edf.setPhysicalDimension(sig, pkt.unit)
+                    edf.setSamplefrequency(sig, pkt.fs)
+                    edf.setLabel(sig, f'{pkt.type}{i}')
+                    data_buffer.append(ch_dat)
+                    sig += 1
+            edf.writeSamples(data_buffer)
+
 
 def _read_block_data(file, shape, format_, format_size):
     sh = np.prod(shape)
@@ -82,26 +128,37 @@ def read_packet(file):
     flow = int.from_bytes(flow, 'big')
     date_ = datetime.fromisoformat(f'{date_[:4]}-{date_[4:6]}-{date_[6:]}T{time_[:2]}:{time_[2:4]}:{time_[4:]}')
     spo2 /= 10.
-    print(date_, flow, hr, spo2, vbat_eeg, vbat_arm)
+    flow = PacketData(np.array([[flow]]), 'flow', 'unit', HEADER_FS)
+    hr = PacketData(np.array([[hr]]), 'hr', 'unit', HEADER_FS)
+    spo2 = PacketData(np.array([[spo2]]), 'spo2', 'unit', HEADER_FS)
+    # print(date_, flow, hr, spo2, vbat_eeg, vbat_arm)
 
     # eeg
-    eeg = _read_block_data(file, (250, 4), 'i', 4).T
-    eeg_giro = _read_block_data(file, (50, 3), 'h', 2).T
-    eeg_acc = _read_block_data(file, (50, 3), 'h', 2).T
-    print(eeg.shape, eeg_giro.shape, eeg_acc.shape)
+    eeg = _read_block_data(file, (250, 4), 'i', 4).T * EEG_EMG_SCALE
+    eeg_giro = _read_block_data(file, (50, 3), 'h', 2).T * GYRO_SCALE
+    eeg_acc = _read_block_data(file, (50, 3), 'h', 2).T * ACC_SCALE
+    eeg = PacketData(eeg, 'eeg', 'uV', EEG_EMG_FS)
+    eeg_giro = PacketData(eeg_giro, 'eeg_giro', 'unit', GYRO_ACC_FS)
+    eeg_acc = PacketData(eeg_acc, 'eeg_acc', 'unit', GYRO_ACC_FS)
+    # print(eeg, eeg_giro, eeg_acc)
 
     # emg
-    emg = _read_block_data(file, (250, 6), 'i', 4).T
-    emg_giro = _read_block_data(file, (50, 3), 'h', 2).T
-    emg_acc = _read_block_data(file, (50, 3), 'h', 2).T
-    print(emg.shape, emg_giro.shape, emg_acc.shape)
+    emg = _read_block_data(file, (250, 6), 'i', 4).T * EEG_EMG_SCALE
+    emg_giro = _read_block_data(file, (50, 3), 'h', 2).T * GYRO_SCALE
+    emg_acc = _read_block_data(file, (50, 3), 'h', 2).T * ACC_SCALE
+    emg = PacketData(emg, 'emg', 'uV', EEG_EMG_FS)
+    emg_giro = PacketData(emg_giro, 'emg_giro', 'unit', GYRO_ACC_FS)
+    emg_acc = PacketData(emg_acc, 'emg_acc', 'unit', GYRO_ACC_FS)
+    # print(emg, emg_giro, emg_acc)
 
     # extra
     sound = _read_block_data(file, (6000, 1), 'h', 2).T
     gsr = _read_block_data(file, (10, 1), 'h', 2).T
-    print(sound.shape, gsr.shape)
+    sound = PacketData(sound, 'sound', 'unit', SOUND_FS)
+    gsr = PacketData(gsr, 'gsr', 'unit', GSR_FS)
+    # print(sound, gsr)
 
-    packet = FlowPacket(date_, flow, hr, spo2,  # vbat_eeg, vbat_arm,
+    packet = FlowPacket(date_, flow, hr, spo2,
                         eeg, eeg_giro, eeg_acc,
                         emg, emg_giro, emg_acc,
                         sound, gsr)
@@ -125,7 +182,8 @@ def read_data(path):
     for file in files:
         with open(file, 'rb') as f:
             data = read_dat_file(f)
-            data.plot()
+            # data.plot()
+            data.to_edf('test.edf')
             exit(12)