py-neuromodulation 0.0.2__py3-none-any.whl → 0.0.3__py3-none-any.whl

packages/tmsi/build/lib/TMSiFileFormats/file_formats/poly5_file_writer.py

@@ -0,0 +1,496 @@
+"""
+(c) 2022 Twente Medical Systems International B.V., Oldenzaal The Netherlands
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+
+#######  #     #   #####   #
+   #     ##   ##  #        
+   #     # # # #  #        #
+   #     #  #  #   #####   #
+   #     #     #        #  #
+   #     #     #        #  #
+   #     #     #  #####    #
+
+/**
+ * @file ${poly5_file_writer.py} 
+ * @brief Poly5 File Writer
+ *
+ */
+
+
+"""
+
+from datetime import datetime
+
+import os
+import threading
+import queue
+import struct
+import time
+import numpy as np
+
+from TMSiSDK.error import TMSiError, TMSiErrorCode
+from TMSiSDK import sample_data_server
+from apex_sdk.device.tmsi_device import TMSiDevice
+from apex_sdk.sample_data_server.sample_data_server import (
+    SampleDataServer as ApexSampleDataServer,
+)
+
+_QUEUE_SIZE = 1000
+
+
+class Poly5Writer:
+    def __init__(self, filename, download=False):
+        self.q_sample_sets = queue.Queue(_QUEUE_SIZE)
+        self.device = None
+
+        fileparts = filename.split(".")
+        if not download:
+            now = datetime.now()
+            filetime = now.strftime("%Y%m%d_%H%M%S")
+            if fileparts[-1] == "poly5" or fileparts[-1] == "Poly5":
+                self.filename = (
+                    ".".join(fileparts[:-1]) + "-" + filetime + ".poly5"
+                )
+            else:
+                self.filename = filename + "-" + filetime + ".poly5"
+        else:
+            if fileparts[-1] == "poly5" or fileparts[-1] == "Poly5":
+                self.filename = filename
+            else:
+                self.filename = filename + ".poly5"
+
+        self._fp = None
+        self._date = None
+
+    def open(self, device):
+        if isinstance(device, TMSiDevice):
+            self.__open_TMSiDevice(device)
+            return
+
+        self.device = device
+        try:
+            self._fp = open(self.filename, "wb")
+            self._date = datetime.now()
+            self._sample_rate = device.config.sample_rate
+            self._num_channels = len(device.channels)
+
+            # Calculate nr of sample-sets within one sample-data-block:
+            # This is the nr of sample-sets in 150 milli-seconds or when the
+            # sample-data-block-size exceeds 64kb the it will become the nr of
+            # sample-sets that fit in 64kb
+            self._num_sample_sets_per_sample_data_block = int(
+                self._sample_rate * 0.15
+            )
+            size_one_sample_set = len(self.device.channels) * 4
+            if (
+                self._num_sample_sets_per_sample_data_block
+                * size_one_sample_set
+            ) > 64000:
+                self._num_sample_sets_per_sample_data_block = int(
+                    64000 / size_one_sample_set
+                )
+
+            # Write poly5-header for thsi measurement
+            Poly5Writer._writeHeader(
+                self._fp,
+                "measurement",
+                device.config.sample_rate,
+                len(device.channels),
+                len(device.channels),
+                0,
+                0,
+                self._date,
+            )
+            for (i, channel) in enumerate(self.device.channels):
+                Poly5Writer._writeSignalDescription(
+                    self._fp, i, channel.name, channel.unit_name
+                )
+
+            fmt = (
+                "f"
+                * self._num_channels
+                * self._num_sample_sets_per_sample_data_block
+            )
+            self.pack_struct = struct.Struct(fmt)
+
+            sample_data_server.registerConsumer(
+                self.device.id, self.q_sample_sets
+            )
+
+            self._sampling_thread = ConsumerThread(
+                self, name="poly5-writer : dev-id-" + str(self.device.id)
+            )
+            self._sampling_thread.start()
+        except OSError as e:
+            print(e)
+            raise TMSiError(TMSiErrorCode.file_writer_error)
+        except:
+            raise TMSiError(TMSiErrorCode.file_writer_error)
+
+    def __open_TMSiDevice(self, device):
+        self.device = device
+        try:
+            self._fp = open(self.filename, "wb")
+            self._date = datetime.now()
+            self._sample_rate = self.device.get_device_sampling_frequency()
+            self._num_channels = self.device.get_num_channels()
+
+            # Calculate nr of sample-sets within one sample-data-block:
+            # This is the nr of sample-sets in 150 milli-seconds or when the
+            # sample-data-block-size exceeds 64kb the it will become the nr of
+            # sample-sets that fit in 64kb
+            self._num_sample_sets_per_sample_data_block = int(
+                self._sample_rate * 0.15
+            )
+            size_one_sample_set = self._num_channels * 4
+            if (
+                self._num_sample_sets_per_sample_data_block
+                * size_one_sample_set
+            ) > 64000:
+                self._num_sample_sets_per_sample_data_block = int(
+                    64000 / size_one_sample_set
+                )
+
+            # Write poly5-header for thsi measurement
+            Poly5Writer._writeHeader(
+                self._fp,
+                "measurement",
+                self._sample_rate,
+                self._num_channels,
+                self._num_channels,
+                0,
+                0,
+                self._date,
+            )
+            for (i, channel) in enumerate(self.device.get_device_channels()):
+                Poly5Writer._writeSignalDescription(
+                    self._fp,
+                    i,
+                    channel.get_channel_name(),
+                    channel.get_channel_unit_name(),
+                )
+
+            fmt = (
+                "f"
+                * self._num_channels
+                * self._num_sample_sets_per_sample_data_block
+            )
+            self.pack_struct = struct.Struct(fmt)
+
+            ApexSampleDataServer().register_consumer(
+                self.device.get_id(), self.q_sample_sets
+            )
+
+            self._sampling_thread = ConsumerThread(
+                self, name="poly5-writer : dev-id-" + str(self.device.get_id())
+            )
+            self._sampling_thread.start()
+        except OSError as e:
+            print(e)
+            raise TMSiError(TMSiErrorCode.file_writer_error)
+        except:
+            raise TMSiError(TMSiErrorCode.file_writer_error)
+
+    def close(self):
+        # print("Poly5Writer-close")
+        self._sampling_thread.stop_sampling()
+
+        if isinstance(self.device, TMSiDevice):
+            ApexSampleDataServer().unregister_consumer(
+                self.device.get_id(), self.q_sample_sets
+            )
+        else:
+            sample_data_server.unregisterConsumer(
+                self.device.id, self.q_sample_sets
+            )
+
+    ## Write header of a poly5 file.
+    #
+    # This function writes the header of a poly5 file to a file.
+    # @param f File object
+    # @param name Name of measurement
+    # @param numSignals Number of signals
+    # @param numSamples Number of samples
+    # @param numDataBlocks Number of data blocks
+    # @param date Date of measurement
+    @staticmethod
+    def _writeHeader(
+        f,
+        name,
+        sample_rate,
+        num_signals,
+        num_samples,
+        num_data_blocks,
+        num_sample_sets_per_sample_data_block,
+        date,
+    ):
+
+        data = struct.pack(
+            "=31sH81phhBHi4xHHHHHHHiHHH64x",
+            b"POLY SAMPLE FILEversion 2.03\r\n\x1a",
+            203,
+            bytes(name, "ascii"),
+            int(sample_rate),
+            int(sample_rate),
+            0,
+            num_signals * 2,
+            num_samples,
+            date.year,
+            date.month,
+            date.day,
+            date.isoweekday() % 7,
+            date.hour,
+            date.minute,
+            date.second,
+            num_data_blocks,
+            num_sample_sets_per_sample_data_block,
+            num_signals * 2 * num_sample_sets_per_sample_data_block * 2,
+            0,
+        )
+        f.write(data)
+
+    ## Write a signal description
+    #
+    # @param f File object
+    # @param index Index of the signal description
+    # @param name Name of the signal (channel)
+    # @param unitname The unit name of the signal
+    @staticmethod
+    def _writeSignalDescription(f, index, name, unit_name):
+        data = struct.pack(
+            "=41p4x11pffffH62x",
+            bytes("(Lo) " + name, "ascii"),
+            bytes(unit_name, "utf-8"),
+            0.0,
+            1000.0,
+            0.0,
+            1000.0,
+            index,
+        )
+        f.write(data)
+
+        data = struct.pack(
+            "=41p4x11pffffH62x",
+            bytes("(Hi) " + name, "ascii"),
+            bytes(unit_name, "utf-8"),
+            0.0,
+            1000.0,
+            0.0,
+            1000.0,
+            index,
+        )
+        f.write(data)
+
+    ## Write a signal block
+    #
+    # @param f File object
+    # @param index Index of the data block
+    # @param date Date of the sample_data block (measurement)
+    # @param signals A list of sample_data, containing NumPy arrays
+    @staticmethod
+    def _writeSignalBlock(
+        f,
+        index,
+        date,
+        sample_sets_block,
+        num_sample_sets_per_sample_data_block,
+        n_chan,
+        pack_struct,
+    ):
+        data = struct.pack(
+            "=i4xHHHHHHH64x",
+            int(index * num_sample_sets_per_sample_data_block),
+            date.year,
+            date.month,
+            date.day,
+            date.isoweekday() % 7,
+            date.hour,
+            date.minute,
+            date.second,
+        )
+        f.write(data)
+
+        sample_sets_block[n_chan - 1 :: n_chan] = sample_sets_block[
+            n_chan - 1 :: n_chan
+        ] % (2**24)
+
+        bin = pack_struct.pack(*sample_sets_block)
+        f.write(bin)
+
+
+class ConsumerThread(threading.Thread):
+    def __init__(self, file_writer, name):
+        super(ConsumerThread, self).__init__()
+        self.name = name
+        self.q_sample_sets = file_writer.q_sample_sets
+        self.sampling = True
+        self._sample_set_block_index = 0
+        self._date = file_writer._date
+        self._fp = file_writer._fp
+        self._sample_rate = file_writer._sample_rate
+        self._num_channels = file_writer._num_channels
+        self._num_sample_sets_per_sample_data_block = (
+            file_writer._num_sample_sets_per_sample_data_block
+        )
+        self._sample_sets_in_block = []
+        self.pack_struct = file_writer.pack_struct
+        self._remaining_samples = np.array([])
+
+    def run(self):
+        # print(self.name, " started")
+
+        while (self.sampling) or (not self.q_sample_sets.empty()):
+            while not self.q_sample_sets.empty():
+                sd = self.q_sample_sets.get()
+                self.q_sample_sets.task_done()
+
+                if self._remaining_samples.size:
+                    samples = np.concatenate(
+                        (self._remaining_samples, sd.samples)
+                    )
+                else:
+                    samples = np.array(sd.samples)
+
+                n_samp = int(len(samples) / sd.num_samples_per_sample_set)
+
+                try:
+                    for i in range(
+                        np.int(
+                            np.floor(
+                                n_samp
+                                / self._num_sample_sets_per_sample_data_block
+                            )
+                        )
+                    ):
+                        self._sample_sets_in_block = samples[
+                            i
+                            * self._num_sample_sets_per_sample_data_block
+                            * sd.num_samples_per_sample_set : (i + 1)
+                            * self._num_sample_sets_per_sample_data_block
+                            * sd.num_samples_per_sample_set
+                        ]
+                        Poly5Writer._writeSignalBlock(
+                            self._fp,
+                            self._sample_set_block_index,
+                            self._date,
+                            self._sample_sets_in_block,
+                            self._num_sample_sets_per_sample_data_block,
+                            self._num_channels,
+                            self.pack_struct,
+                        )
+                        self._sample_set_block_index += 1
+
+                        if not (self._sample_set_block_index % 20):
+                            # Go back to start and rewrite header
+                            self._fp.seek(0)
+                            Poly5Writer._writeHeader(
+                                self._fp,
+                                "measurement",
+                                self._sample_rate,
+                                self._num_channels,
+                                self._sample_set_block_index
+                                * self._num_sample_sets_per_sample_data_block,
+                                self._sample_set_block_index,
+                                self._num_sample_sets_per_sample_data_block,
+                                self._date,
+                            )
+
+                            # Flush all data from buffers to the file
+                            self._fp.flush()
+                            os.fsync(self._fp.fileno())
+
+                            # Go back to end of file
+                            self._fp.seek(0, os.SEEK_END)
+
+                    i = np.int(
+                        np.floor(
+                            n_samp / self._num_sample_sets_per_sample_data_block
+                        )
+                    )
+                    ind = np.arange(
+                        i
+                        * self._num_sample_sets_per_sample_data_block
+                        * sd.num_samples_per_sample_set,
+                        n_samp * sd.num_samples_per_sample_set,
+                    )
+                    if ind.any:
+                        self._remaining_samples = samples[ind]
+                    else:
+                        self._remaining_samples = np.array([])
+
+                except:
+                    raise TMSiError(TMSiErrorCode.file_writer_error)
+
+            time.sleep(0.01)
+
+        while self._remaining_samples.any():
+            # Last data block is omitted from saving, to prevent an incomplete data block being part of the Poly5 file
+            # This would result in 0s at the end of the file
+            if (
+                np.shape(self._remaining_samples)[0]
+                < self._num_sample_sets_per_sample_data_block
+                * self._num_channels
+            ):
+                self._remaining_samples = np.array([])
+
+            else:
+                self._sample_sets_in_block = self._remaining_samples[
+                    : self._num_sample_sets_per_sample_data_block
+                    * self._num_channels
+                ]
+
+                self._remaining_samples = self._remaining_samples[
+                    self._num_sample_sets_per_sample_data_block
+                    * self._num_channels :
+                ]
+
+                Poly5Writer._writeSignalBlock(
+                    self._fp,
+                    self._sample_set_block_index,
+                    self._date,
+                    self._sample_sets_in_block,
+                    self._num_sample_sets_per_sample_data_block,
+                    self._num_channels,
+                    self.pack_struct,
+                )
+                self._sample_set_block_index += 1
+
+        # Go back to start and rewrite header
+        self._fp.seek(0)
+        Poly5Writer._writeHeader(
+            self._fp,
+            "measurement",
+            self._sample_rate,
+            self._num_channels,
+            self._sample_set_block_index
+            * self._num_sample_sets_per_sample_data_block,
+            self._sample_set_block_index,
+            self._num_sample_sets_per_sample_data_block,
+            self._date,
+        )
+
+        # Flush all data from buffers to the file
+        self._fp.flush()
+        os.fsync(self._fp.fileno())
+
+        # Go back to end of file
+        self._fp.seek(0, os.SEEK_END)
+
+        # print(self.name, " ready, closing file")
+        self._fp.close()
+        return
+
+    def stop_sampling(self):
+        # print(self.name, " stop sampling")
+        self.sampling = False