diff --git a/.gitignore b/.gitignore index 11aa7d7..8d900a6 100644 --- a/.gitignore +++ b/.gitignore @@ -9,6 +9,11 @@ __pycache__/ # bio extensions *.bio +# vs code settings +*.env +.vscode/ + + # Distribution / packaging .idea .Python diff --git a/biosiglive/__init__.py b/biosiglive/__init__.py index f14b3fb..3373df1 100644 --- a/biosiglive/__init__.py +++ b/biosiglive/__init__.py @@ -7,7 +7,7 @@ from .interfaces.param import Param, Device, MarkerSet -from .file_io.save_and_load import load, save +from .file_io.save_and_load import load, save, compress from .processing.data_processing import RealTimeProcessing, OfflineProcessing, GenericProcessing from .processing.msk_functions import MskFunctions @@ -16,7 +16,12 @@ from .streaming.client import Client, Message from .streaming.server import Server from .streaming.stream_data import StreamData +from .streaming.async_client import AsyncTCPClient +from .streaming.async_server import AsyncTCPServer from .enums import * -__version__ = "2.0.0" +from .interfaces.trigno_sdk.sdk_client import TrignoSDKClient + + +__version__ = "2.0.4" diff --git a/biosiglive/enums.py b/biosiglive/enums.py index 5c103fa..edd05db 100644 --- a/biosiglive/enums.py +++ b/biosiglive/enums.py @@ -1,6 +1,7 @@ """ Groups the different enums used in the program. it is a good place to start to check what's available. """ + from enum import Enum @@ -22,6 +23,7 @@ class DeviceType(Enum): Emg = "emg" Imu = "imu" + DelsysGogniometer = "delsys_gogniometer" Generic = "generic" ForcePlate = "force_plate" diff --git a/biosiglive/file_io/save_and_load.py b/biosiglive/file_io/save_and_load.py index f15cac1..20e41e0 100644 --- a/biosiglive/file_io/save_and_load.py +++ b/biosiglive/file_io/save_and_load.py @@ -1,9 +1,63 @@ +import os.path + import pickle -import numpy as np from pathlib import Path +from ..streaming.utils import dic_merger +import pickletools +import gzip +from multiprocessing import Pool + + +def _load_and_compress(origin_file, target_file=None): + data = load(origin_file) + if not target_file: + target_file = origin_file.split(".")[:-1][0] + ".bio.gzip" + if Path(target_file).suffix != ".gzip": + target_file = target_file + ".gzip" + f = gzip.open(target_file, "wb") + data_pick = pickletools.optimize(pickle.dumps(data, pickle.HIGHEST_PROTOCOL)) + f.write(data_pick) + f.close() + + +def compress(origin_file, target_file=None, multi_proc=False): + if target_file and len(target_file) != len(origin_file): + raise ValueError("The number of target files must be equal to the number of origin files.") + if multi_proc: + pool = Pool(processes=os.cpu_count()) + pool.map(_load_and_compress, origin_file) + else: + for i in range(len(origin_file)): + target_file_tmp = target_file[i] if target_file else None + _load_and_compress(origin_file[i], target_file_tmp) + + +def _safe_rename_file(file_name) -> str: + """This function checks if the file extension is an integer. + Parameters + ---------- + file_name : str + The name to the file. + + Returns + ------- + str + new name of the file. + """ + i = 0 + while file_name[-(i + 1)].isdigit(): + i += 1 + old_value = int(file_name[-i:]) if i > 0 else None + if old_value: + if file_name[-(i + 1)] == "_" or file_name[-(i + 1)] == "-": + i += 1 + return file_name[:-i] + "_" + str(old_value + 1) + else: + return file_name + "_1" -def save(data_dict, data_path): + +def save(data_dict, data_path, add_data=False, safe=True, compress=False): """This function adds data to a pickle file. It not open the file, but appends the data to the end, so it's fast. Parameters @@ -12,18 +66,44 @@ def save(data_dict, data_path): The data to be added to the file. data_path : str The path to the file. The file must exist. + add_data : bool, optional + If True, the data are added to the file. If False, the file is overwritten. The default is False. + safe : bool, optional + If True, the data are saved in a new file. The default is False. + compress : bool, optional + If True, the data are compressed. The default is False. """ - if Path(data_path).suffix != ".bio": - if Path(data_path).suffix == "": + data_path_object = Path(data_path) + if data_path_object.suffix != ".bio": + if data_path_object.suffix == "": data_path += ".bio" else: raise ValueError("The file must be a .bio file.") - with open(data_path, "ab") as outf: - pickle.dump(data_dict, outf, pickle.HIGHEST_PROTOCOL) + if not add_data: + if safe and os.path.isfile(data_path): + file_name = _safe_rename_file(data_path) + data_path = data_path_object.parent / (file_name + data_path_object.suffix) + print( + f"The file {data_path_object.name} already exists. The data will be saved in {data_path.name}." + f" To avoid this message, remove the safe option." + ) + if not safe and os.path.isfile(data_path): + os.remove(data_path) + + if compress: + if add_data: + raise ValueError("The data cannot be added to a compressed file.") + data_path = data_path + ".gzip" + f = gzip.open(data_path, "wb") + data_pick = pickletools.optimize(pickle.dumps(data_dict, pickle.HIGHEST_PROTOCOL)) + f.write(data_pick) + f.close() + else: + with open(data_path, "ab") as outf: + pickle.dump(data_dict, outf, pickle.HIGHEST_PROTOCOL) -# TODO add dict merger -def load(filename, number_of_line=None): +def load_old(filename, number_of_line=None, merge=True): """This function reads data from a pickle file to concatenate them into one dictionary. Parameters @@ -32,6 +112,8 @@ def load(filename, number_of_line=None): The path to the file. number_of_line : int The number of lines to read. If None, all lines are read. + merge : bool + If True, the data are merged into one dictionary. If False, the data are returned as a list of dictionaries. Returns ------- @@ -39,32 +121,106 @@ def load(filename, number_of_line=None): The data read from the file. """ - if Path(filename).suffix != ".bio": - raise ValueError("The file must be a .bio file.") - data = {} + with_gzip = False + + if Path(filename).suffix == ".bio": + with_gzip = False + elif Path(filename).suffix == ".gzip": + with_gzip = True + # else: + # raise ValueError("The file must be a .bio or a .bio.gzip file.") + data = None if merge else [] limit = 2 if not number_of_line else number_of_line - with open(filename, "rb") as file: + if with_gzip: + with gzip.open(filename, "rb") as file: + count = 0 + while count < limit: + try: + data_tmp = pickle.load(file) + if not merge: + data.append(data_tmp) + else: + if not data: + data = data_tmp + else: + data = dic_merger(data, data_tmp) + if number_of_line: + count += 1 + else: + count = 1 + except EOFError: + break + else: + with open(filename, "rb") as file: + count = 0 + while count < limit: + try: + data_tmp = pickle.load(file) + if not merge: + data.append(data_tmp) + else: + if not data: + data = data_tmp + else: + data = dic_merger(data, data_tmp) + if number_of_line: + count += 1 + else: + count = 1 + except EOFError: + break + return data + + +def load(filename, number_of_line=None, merge=True): + """This function reads data from a pickle file to concatenate them into one dictionary. + + Parameters + ---------- + filename : str + The path to the file. + number_of_line : int + The number of lines to read. If None, all lines are read. + merge : bool + If True, the data are merged into one dictionary. If False, the data are returned as a list of dictionaries. + + Returns + ------- + data : dict + The data read from the file. + + """ + if not Path(filename).is_file(): + raise FileNotFoundError(f"The file {filename} does not exist.") + + with_gzip = filename.endswith(".gzip") + data = None if merge else [] + limit = number_of_line if number_of_line else float("inf") + + try: + return _read_all_lines(filename, limit, data, with_gzip, merge) + except Exception as e: + raise RuntimeError(f"An error occurred while loading the file: {e}") + + +def _read_all_lines(filename, limit=float("inf"), data=(), with_gzip=False, merge=True): + file_open = gzip.open if with_gzip else open + with file_open(filename, "rb") as file: count = 0 while count < limit: try: data_tmp = pickle.load(file) - for key in data_tmp.keys(): - if key in data.keys(): - if isinstance(data[key], list) is True: - data[key].append(data_tmp[key]) - else: - data[key] = np.append(data[key], data_tmp[key], axis=len(data[key].shape) - 1) - else: - if isinstance(data_tmp[key], (int, float, str, dict)) is True: - data[key] = [data_tmp[key]] - elif isinstance(data_tmp[key], list) is True: - data[key] = [data_tmp[key]] - else: - data[key] = data_tmp[key] - if number_of_line: - count += 1 + if not merge: + data.append(data_tmp) + else: + data = dic_merger(data, data_tmp) if data else data_tmp + count += 1 + except Exception as e: + if str(e) == "Ran out of input": + # print(f"{count} lines read from the file") + pass else: - count = 1 - except EOFError: + print(f"An error occurred while reading the file: {e} at line {count}") break + return data diff --git a/biosiglive/gui/plot.py b/biosiglive/gui/plot.py index 334320f..7a3025f 100644 --- a/biosiglive/gui/plot.py +++ b/biosiglive/gui/plot.py @@ -1,21 +1,52 @@ """ This file contains all the plot functions to plot the data in live or offline. """ + try: import pyqtgraph as pg import pyqtgraph.opengl as gl - from PyQt5.QtWidgets import QProgressBar +except ModuleNotFoundError: + pass + +import time + +try: + from PyQt5.QtCore import QTimer + from PyQt5.QtWidgets import QMainWindow, QApplication +except ModuleNotFoundError: + pass + +try: + from PyQt6.QtWidgets import QProgressBar except ModuleNotFoundError: pass import numpy as np from typing import Union -import matplotlib.pyplot as plt +import sys + +try: + import matplotlib.pyplot as plt +except ModuleNotFoundError: + pass + from math import ceil from ..enums import PlotType -import time +from ..streaming.utils import CircularBuffer + +pyqt_color_list = [ + (31, 119, 180), # muted blue + (255, 127, 14), # soft orange + (44, 160, 44), # soft green + (214, 39, 40), # soft red + (148, 103, 189), # purple + (140, 86, 75), # brown + (227, 119, 194), # pink + (127, 127, 127), # gray + (188, 189, 34), # olive + (23, 190, 207), # cyan +] -# TODO: Add plot class to be able to plot several curves in the same plot. And do several plot in the same app. class LivePlot: def __init__( self, @@ -23,7 +54,7 @@ def __init__( plot_type: Union[PlotType, str] = PlotType.Curve, name: str = None, channel_names: list = None, - rate: int = None, + rate: int = 60, ): """ Initialize the plot class. @@ -49,10 +80,18 @@ def __init__( if len(channel_names) != nb_subplots: raise ValueError("The number of subplots is not equal to the number of channel names") + self.rate = rate + + # + # self.app = pg.mkQApp("Curve_plot") + # self.app = QApplication([]) + + # self.timer = QTimer() + # self.timer.timeout.connect(self._update_plot) + self.channel_names = channel_names self.figure_name = name self.nb_subplot = nb_subplots - self.rate = rate self.layout = None self.app = None self.viz = None @@ -72,6 +111,7 @@ def __init__( def init( self, plot_windows: Union[int, list] = None, + create_app: bool = True, **kwargs, ): """ @@ -82,7 +122,9 @@ def init( plot_windows: Union[int, list] The number of frames ti plot. If is a list, the number of frames to plot for each subplot. """ - self.plot_buffer = [None] * self.nb_subplot + self.create_app = create_app + + # self.plot_buffer = [None] * self.nb_subplot if isinstance(plot_windows, int): plot_windows = [plot_windows] * self.nb_subplot self.plot_windows = plot_windows @@ -101,6 +143,27 @@ def init( else: raise ValueError(f"The plot type ({self.plot_type}) is not supported.") + # self.timer.start(int((1/self.rate) * 1000)) + self.win.show() + # self.app.exec() + + def _update_plot(self, **kwargs): + + if not self.once_update: + return + + if self.plot_type == PlotType.ProgressBar: + self._update_progress_bar(self.get_unflatten(self.plot_buffer.get())) + elif self.plot_type == PlotType.Curve: + self._update_curve(self.get_unflatten(self.plot_buffer.get())) + elif self.plot_type == PlotType.Skeleton: + self._update_skeleton(self.get_unflatten(self.plot_buffer.get()), self.viz) + elif self.plot_type == PlotType.Scatter3D: + self._update_3d_scatter(self.get_unflatten(self.plot_buffer.get()), **kwargs) + else: + raise ValueError(f"The plot type ({self.plot_type}) is not supported.") + self.app.processEvents() + def update(self, data: Union[np.ndarray, list], **kwargs): """ This function is used to update the qt app. @@ -110,7 +173,6 @@ def update(self, data: Union[np.ndarray, list], **kwargs): data: Union[np.ndarray, list] The data to plot. If it is a list, the data to plot for each subplot. """ - update = True if self.plot_type != PlotType.Scatter3D and self.plot_type != PlotType.Skeleton: if isinstance(data, list): if len(data) != self.nb_subplot: @@ -130,40 +192,34 @@ def update(self, data: Union[np.ndarray, list], **kwargs): data.append(d[np.newaxis, :]) if self.plot_windows: - for i in range(self.nb_subplot): - if self.plot_buffer[i] is None: - self.plot_buffer[i] = data[i][..., -self.plot_windows[i] :] - if self.plot_buffer[i].shape[1] < self.plot_windows[i]: - size = self.plot_windows[i] - self.plot_buffer[i].shape[1] - self.plot_buffer[i] = np.append( - np.zeros((self.plot_buffer[i].shape[0], size)), self.plot_buffer[i], axis=-1 - ) - elif self.plot_buffer[i].shape[1] < self.plot_windows[i]: - self.plot_buffer[i] = np.append(self.plot_buffer[i], data[i], axis=-1) - elif self.plot_buffer[i].shape[1] >= self.plot_windows[i]: - size = data[i].shape[1] - self.plot_buffer[i] = np.append(self.plot_buffer[i][..., size:], data[i], axis=-1) - data = self.plot_buffer + if self.plot_buffer is None: + flat_data = self.get_flatten(data) + self.plot_buffer = CircularBuffer(flat_data.shape[0], self.plot_windows[0]) + self.plot_buffer.append(flat_data) + else: + self.plot_buffer.append(self.get_flatten(data)) + + update = True if self.rate and self.once_update: plot_time = time.time() - self.last_plot if plot_time != 0 and 1 / plot_time > self.rate: update = False else: update = True + if update: self.once_update = True - if self.plot_type == PlotType.ProgressBar: - self._update_progress_bar(data) - elif self.plot_type == PlotType.Curve: - self._update_curve(data) - elif self.plot_type == PlotType.Skeleton: - self._update_skeleton(data, self.viz) - elif self.plot_type == PlotType.Scatter3D: - self._update_3d_scatter(data, **kwargs) - else: - raise ValueError(f"The plot type ({self.plot_type}) is not supported.") + self._update_plot(**kwargs) self.last_plot = time.time() + def get_flatten(self, data): + sizes = np.array([len(d) for d in data]) + self._offsets = np.concatenate(([0], np.cumsum(sizes))) + return np.vstack(data) + + def get_unflatten(self, data): + return [data[self._offsets[i] : self._offsets[i + 1]] for i in range(len(self._offsets) - 1)] + def _init_curve( self, figure_name: str = "Figure", @@ -198,6 +254,7 @@ def _init_curve( self.app = pg.mkQApp("Curve_plot") pg.setConfigOption("background", "w") pg.setConfigOption("foreground", "k") + self.win = pg.GraphicsLayoutWidget(show=True) self.win.setWindowTitle(figure_name) nb_line = 4 @@ -244,8 +301,11 @@ def _init_curve( line_count = 0 self.plots.append(self.win.addPlot(title=subplot_labels[subplot])) self.plots[-1].setDownsampling(mode="peak") - self.plots[-1].setClipToView(False) - self.curves.append(self.plots[-1].plot([], pen=colors[subplot], name="Blue curve")) + self.plots[-1].setClipToView(True) + self.plots[-1].enableAutoRange(False) + # self.curves.append(self.plots[-1].multiDataPlot(x=[], y=[], pen=colors[subplot], name="Blue curve")) + + # self.curves.append(self.plots[-1].plot([], pen=colors[subplot], name="Blue curve")) self.plots[-1].setLabel("bottom", x_labels[subplot]) self.plots[-1].setLabel("left", y_labels[subplot]) self.plots[-1].showGrid(x=grid, y=grid) @@ -275,7 +335,7 @@ def _init_progress_bar( # --- Progress bar graph --- # if isinstance(unit, str): self.unit = [unit] * nb_subplot - self.layout, self.app = self._init_layout(figure_name, resize=self.resize, move=self.move) + self._init_layout(figure_name, resize=self.resize, move=self.move) row_count = 0 if bar_graph_max_value is None: bar_graph_max_value = [100] * nb_subplot @@ -351,7 +411,6 @@ def _update_3d_scatter( raise ValueError("The number of size is not equal to the number of data.") for plot in self.plots: plot.setData(pos=data, color=colors, size=size) - self.app.processEvents() def _update_curve(self, data: list): """ @@ -362,17 +421,38 @@ def _update_curve(self, data: list): data: list The data to plot. """ - if len(data) != len(self.curves): - raise ValueError( - f"The number of data ({len(data)}) is different from the number of curves ({len(self.curves)})." - ) + # if len(data) != len(self.curves): + # self._add_curve() + # raise ValueError( + # f"The number of data ({len(data)}) is different from the number of curves ({len(self.curves)})." + # ) + if len(self.curves) == 0: + for d, da in enumerate(data): + if da.shape[0] == 1: + self.curves.append( + self.plots[d].plot( + da[0, :], pen=pg.mkPen(color=(0, 128, 232)), downsample=5, autoDownsample=True + ) + ) + else: + self.curves.append( + self.plots[d].multiDataPlot( + x=np.arange(da.shape[1]), + y=da, + pen=[pg.mkPen(color=pyqt_color_list[k]) for k in range(da.shape[0])], + downsample=[5 for k in range(da.shape[0])], + autoDownsample=[True for k in range(da.shape[0])], + ) + ) + for i in range(len(data)): - if self.ptr[i] == 0: - self.size_to_append[i] = data[i].shape[1] - self.ptr[i] += self.size_to_append[i] * 2 - self.curves[i].setData(data[i][0, :]) - # self.curves[i].setPos(self.ptr[i], 0) - self.app.processEvents() + if data[i].shape[0] == 1: + self.curves[i].setData(y=data[i][0, :], downsample=5, autoDownsample=True) + else: + [ + self.curves[i][j].setData(y=data[i][j], downsample=5, autoDownsample=True) + for j in range(len(self.curves[i])) + ] def _update_progress_bar(self, data: list): """ @@ -395,7 +475,6 @@ def _update_progress_bar(self, data: list): self.plots[i].setValue(int(value)) name = self.channel_names[i] if self.channel_names else f"plot_{i}" self.plots[i].setFormat(f"{name}: {int(value)} {self.unit[i]}") - self.app.processEvents() @staticmethod def _update_skeleton(data: list, viz): @@ -429,8 +508,7 @@ def _init_skeleton(**kwargs): plot = bioviz.Viz(**kwargs) return plot - @staticmethod - def _init_layout(figure_name: str = "Figure", resize: tuple = (400, 400), move: tuple = (0, 0)): + def _init_layout(self, figure_name: str = "Figure", resize: tuple = (400, 400), move: tuple = (0, 0)): """ This function is used to initialize the qt app layout. @@ -451,11 +529,11 @@ def _init_layout(figure_name: str = "Figure", resize: tuple = (400, 400), move: The qt app. """ - app = pg.mkQApp(figure_name) - layout = pg.LayoutWidget() - layout.resize(resize[0], resize[1]) - layout.move(move[0], move[1]) - return layout, app + self.graph_widget = pg.mkQApp(figure_name) + self.layout = pg.LayoutWidget() + self.layout.resize(resize[0], resize[1]) + self.layout.move(move[0], move[1]) + self.setCentralWidget(self.layout) def disconnect(self): self.app.disconnect() diff --git a/biosiglive/interfaces/generic_interface.py b/biosiglive/interfaces/generic_interface.py index 580f971..1399af6 100644 --- a/biosiglive/interfaces/generic_interface.py +++ b/biosiglive/interfaces/generic_interface.py @@ -1,6 +1,7 @@ """ This file contains a generic interface class to use for any new implemented class. """ + from .param import * from typing import Union from ..enums import DeviceType, InverseKinematicsMethods, InterfaceType @@ -12,7 +13,7 @@ class GenericInterface: """ def __init__( - self, ip: str = "127.0.0.1", system_rate: float = 100, interface_type: Union[InterfaceType, str] = None + self, ip: str = "127.0.0.1", system_rate: float = None, interface_type: Union[InterfaceType, str] = None ): """ Initialize the generic class. @@ -45,8 +46,7 @@ def _add_device( nb_channels: int, device_type: Union[DeviceType, str] = DeviceType.Emg, name: str = None, - rate: float = 2000, - device_range: tuple = None, + rate: float = None, processing_method: Union[RealTimeProcessingMethod, OfflineProcessingMethod] = None, **kwargs, ): @@ -75,7 +75,6 @@ def _add_device( raise ValueError("The type of the device is not valid.") device_type = DeviceType(device_type) device_tmp = Device(device_type, nb_channels, name, rate, self.system_rate) - device_tmp.device_range = device_range if device_range else (0, nb_channels) device_tmp.interface = self.interface_type device_tmp.processing_method = processing_method device_tmp.processing_method_kwargs = kwargs diff --git a/biosiglive/interfaces/param.py b/biosiglive/interfaces/param.py index 31b96f6..e1053fa 100644 --- a/biosiglive/interfaces/param.py +++ b/biosiglive/interfaces/param.py @@ -1,6 +1,7 @@ """ This file contains the Parameter class that define the device and markers classes. """ + from math import ceil from ..enums import DeviceType, MarkerType, InverseKinematicsMethods, RealTimeProcessingMethod, OfflineProcessingMethod from ..processing.data_processing import RealTimeProcessing, OfflineProcessing, GenericProcessing @@ -13,8 +14,8 @@ class Param: def __init__( self, nb_channels: int, + rate: float, name: str = None, - rate: float = None, system_rate: float = 100, data_window: int = None, ): @@ -38,10 +39,10 @@ def __init__( self.name = name self.rate = rate self.system_rate = system_rate - self.sample = ceil(rate / self.system_rate) + self.sample = None if rate is None else ceil(rate / self.system_rate) self.range = None self.raw_data = [] - self.data_window = data_window if data_window else int(rate) + self.data_window = data_window if data_window else int(self.rate * 2) self.new_data = None def append_data(self, new_data: np.ndarray): @@ -71,7 +72,7 @@ def __init__( device_type: DeviceType = DeviceType.Emg, nb_channels: int = 1, name: str = None, - rate: float = 2000, + rate: float = None, system_rate: float = 100, channel_names: Union[list, str] = None, ): @@ -94,7 +95,7 @@ def __init__( channel_names: list Name of the channels of the device. """ - super().__init__(nb_channels, name, rate, system_rate) + super().__init__(nb_channels, rate, name, system_rate) if isinstance(channel_names, str): channel_names = [channel_names] if channel_names: @@ -189,7 +190,7 @@ def _init_processing_method(self): elif self.processing_method == RealTimeProcessingMethod.Custom: self.processing_function = RealTimeProcessing(self.rate, self.processing_window).custom_processing - def _check_if_has_changed(self, method: GenericProcessing(), kwargs: dict) -> bool: + def _check_if_has_changed(self, method: GenericProcessing, kwargs: dict) -> bool: """ Check if the processing method has changed. @@ -267,7 +268,7 @@ def __init__( Unit of the marker set. """ marker_type = MarkerType.Unlabeled if unlabeled else MarkerType.Labeled - super(MarkerSet, self).__init__(nb_channels, name, rate, system_rate) + super(MarkerSet, self).__init__(nb_channels, rate, name, system_rate) if isinstance(marker_names, str): marker_names = [marker_names] if marker_names: diff --git a/biosiglive/interfaces/pytrigno_interface.py b/biosiglive/interfaces/pytrigno_interface.py index bdcbc3a..49b3bf7 100644 --- a/biosiglive/interfaces/pytrigno_interface.py +++ b/biosiglive/interfaces/pytrigno_interface.py @@ -3,10 +3,11 @@ from ..enums import DeviceType, InterfaceType, RealTimeProcessingMethod, OfflineProcessingMethod from typing import Union -try: - import pytrigno -except ModuleNotFoundError: - pass +# try: +# import pytrigno +# except ModuleNotFoundError: +# pass +from .trigno_sdk.sdk_client import TrignoSDKClient class PytrignoClient(GenericInterface): @@ -14,7 +15,7 @@ class PytrignoClient(GenericInterface): Class to wrap the Trigno community SDK. """ - def __init__(self, system_rate=100, ip: str = "127.0.0.1", init_now: bool = True): + def __init__(self, system_rate=74.074074, ip: str = "127.0.0.1", init_now: bool = True): """ Initialize the interface. @@ -39,6 +40,12 @@ def __init__(self, system_rate=100, ip: str = "127.0.0.1", init_now: bool = True self.is_frame = False self.is_initialized = False self.init_now = init_now + self.sdk_client = None + # if system_rate != 74.074074: + # raise ValueError("System rate can not be changed for pytrigno interface for now." \ + # " 74 Hz mean that data will refresh every 13.5 ms but the EMG and IMU data are sampled at their own rate.") + if self.init_now: + self.sdk_client = TrignoSDKClient(host=ip, init_sensors=False, stream_rate=system_rate) def add_device( self, @@ -47,7 +54,6 @@ def add_device( data_buffer_size: int = None, name: str = None, rate: float = 2000, - device_range: tuple = None, processing_method: Union[RealTimeProcessingMethod, OfflineProcessingMethod] = None, **process_kwargs, ): @@ -57,7 +63,7 @@ def add_device( Parameters ---------- nb_channels: int - Number of channels of the device. + Number of channels of the device will be remove in futur version. device_type: Union[DeviceType, str] Type of the device. data_buffer_size: int @@ -73,11 +79,7 @@ def add_device( **process_kwargs Keyword arguments for the processing method. """ - if not device_range: - device_range = (0, nb_channels - 1) - device_tmp = self._add_device( - nb_channels, device_type, name, rate, device_range, processing_method, **process_kwargs - ) + device_tmp = self._add_device(nb_channels, device_type, name, rate, processing_method, **process_kwargs) device_tmp.interface = self.interface_type device_tmp.data_windows = data_buffer_size self.devices.append(device_tmp) @@ -85,8 +87,13 @@ def add_device( if device_type not in [t.value for t in DeviceType]: raise ValueError("The type of the device is not valid.") device_type = DeviceType(device_type) - if device_type != DeviceType.Emg and device_type != DeviceType.Imu: - raise RuntimeError("Device type must be 'emg' or 'imu' with pytrigno.") + if ( + device_type != DeviceType.Emg + and device_type != DeviceType.Imu + and device_type != DeviceType.DelsysGogniometer + ): + raise RuntimeError("Device type must be 'emg', 'delsys_gogniometer' or 'imu' with pytrigno.") + if self.init_now: self.init_client() @@ -130,22 +137,12 @@ def get_device_data( for device in devices: if get_frame: - if device.device_type == DeviceType.Emg: - count = 0 - for device_tmp in devices: - if device_tmp == device: - break - if device.device_type == DeviceType.Emg: - count += 1 - device.new_data = self.emg_client[count].read() - else: - count = 0 - for device_tmp in devices: - if device_tmp == device: - break - if device.device_type == DeviceType.Imu: - count += 1 - device.new_data = self.imu_client[count].read() + if device.device_type == DeviceType.Emg or device.device_type == DeviceType.DelsysGogniometer: + queue_name = [key for key in self.sdk_client.all_queue if "emg" in key][0] + elif device.device_type == DeviceType.Imu or device.device_type == DeviceType.DelsysGogniometer: + queue_name = [key for key in self.sdk_client.all_queue if "aux" in key][0] + device.new_data, _ = self.sdk_client.all_queue[queue_name].get() + if channel_idx: device_data = np.ndarray((len(channel_idx), device.new_data.shape[1])) for i, idx in enumerate(channel_idx): @@ -163,28 +160,31 @@ def get_frame(self): Get a frame from the interface. This function is used to get data from the interface. """ if not self.is_initialized: - raise RuntimeError("Client is not initialized. Please call init_client() first.") + self.init_client() + # raise RuntimeError("Client is not initialized. Please call init_client() first.") self.get_device_data(get_frame=True) return True + def check_rate_devices(self): + for device in self.devices: + if "emg" in device.device_type.value: + device.rate = self.sdk_client.get_emg_streaming_rate() + if "imu" in device.device_type.value: + device.rate = self.sdk_client.get_aux_streaming_rate() + def init_client(self): """ Initialize the client if it's not already done. This function has to be called before getting a frame. """ self.is_initialized = True - for d, device in enumerate(self.devices): - if device.device_type == DeviceType.Emg: - self.emg_client.append( - pytrigno.TrignoEMG( - channel_range=device.device_range, samples_per_read=device.sample, host=self.address - ) - ) - self.emg_client[-1].start() - elif device.device_type == DeviceType.Imu: - imu_range = (device.device_range[0] * 9, device.device_range[1] * 9) - self.imu_client.append( - pytrigno.TrignoIM(channel_range=imu_range, samples_per_read=device.sample, host=self.address) + device_types = np.unique(np.array([device.device_type.value for device in self.devices])) + self.sdk_client.initialize_sensors(device_types) + self.check_rate_devices() + if self.sdk_client._threads_to_run["avanti_emg"] and self.sdk_client._threads_to_run["legacy_emg"]: + raise ( + RuntimeError( + "Both avanti and legacy emg type are paired. " + "It is not possible to use both at the same time in biosiglive for now." ) - self.imu_client[-1].start() - else: - raise RuntimeError("Device type must be 'emg' or 'imu' with pytrigno.") + ) + self.sdk_client.start_streaming() diff --git a/biosiglive/interfaces/tcp_interface.py b/biosiglive/interfaces/tcp_interface.py index c31784e..6c41fa6 100644 --- a/biosiglive/interfaces/tcp_interface.py +++ b/biosiglive/interfaces/tcp_interface.py @@ -1,6 +1,7 @@ """ This file contains a wrapper to use a tcp client more easily. """ + from ..streaming.client import Client, Message from .generic_interface import GenericInterface from ..enums import ( @@ -37,6 +38,7 @@ def __init__(self, ip: str = "127.0.0.1", port: int = 801, client_type: str = "T self.devices = [] self.imu = [] self.marker_sets = [] + self.system_rate = read_frequency self.read_frequency = read_frequency self.ip = ip self.port = port @@ -52,7 +54,6 @@ def add_device( command_name: Union[str, list] = "", name: str = None, rate: float = 2000, - device_range: tuple = None, processing_method: Union[RealTimeProcessingMethod, OfflineProcessingMethod] = None, **process_kwargs, ): @@ -71,16 +72,12 @@ def add_device( Name of the device. rate: float Frequency of the device. - device_range: tuple - Range of the device. processing_method: Union[RealTimeProcessingMethod, OfflineProcessingMethod] Method to use to process the data. process_kwargs: dict Keyword arguments for the processing method. """ - device_tmp = self._add_device( - nb_channels, device_type, name, rate, device_range, processing_method, **process_kwargs - ) + device_tmp = self._add_device(nb_channels, device_type, name, rate, processing_method, **process_kwargs) device_tmp.interface = self.interface_type self.devices.append(device_tmp) self.device_cmd_names.append(command_name) diff --git a/biosiglive/interfaces/trigno_sdk/__init__.py b/biosiglive/interfaces/trigno_sdk/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/biosiglive/interfaces/trigno_sdk/avanti_mode_enum.py b/biosiglive/interfaces/trigno_sdk/avanti_mode_enum.py new file mode 100644 index 0000000..6d69a8c --- /dev/null +++ b/biosiglive/interfaces/trigno_sdk/avanti_mode_enum.py @@ -0,0 +1,140 @@ +from enum import Enum + + +class EMGAvantiMode(Enum): + EMG_ACC_2G = 0 + EMG_ACC_4G = 1 + EMG_ACC_8G = 2 + EMG_ACC_16G = 3 + EMG_GYRO_250DPS = 4 + EMG_GYRO_500DPS = 5 + EMG_GYRO_1000DPS = 6 + EMG_GYRO_2000DPS = 7 + EMG_IMU_2G_250DPS = 8 + EMG_IMU_4G_250DPS = 9 + EMG_IMU_8G_250DPS = 10 + EMG_IMU_16G_250DPS = 11 + EMG_IMU_2G_500DPS = 12 + EMG_IMU_4G_500DPS = 13 + EMG_IMU_8G_500DPS = 14 + EMG_IMU_16G_500DPS = 15 + EMG_IMU_2G_1000DPS = 16 + EMG_IMU_4G_1000DPS = 17 + EMG_IMU_8G_1000DPS = 18 + EMG_IMU_16G_1000DPS = 19 + EMG_IMU_2G_2000DPS = 20 + EMG_IMU_4G_2000DPS = 21 + EMG_IMU_8G_2000DPS = 22 + EMG_IMU_16G_2000DPS = 23 + EMG_ORIENTATION_5CH = 39 + EMG_ONLY_2148HZ = 40 + EMG_ACC_74HZ_2G = 42 + EMG_ACC_74HZ_4G = 43 + EMG_ACC_74HZ_8G = 44 + EMG_ACC_74HZ_16G = 45 + EMG_GYRO_74HZ_250DPS = 46 + EMG_GYRO_74HZ_500DPS = 47 + EMG_GYRO_74HZ_1000DPS = 48 + EMG_GYRO_74HZ_2000DPS = 49 + EMG_ACC_GYRO_148HZ_2G_250DPS = 50 + EMG_ACC_GYRO_148HZ_4G_250DPS = 51 + EMG_ACC_GYRO_148HZ_8G_250DPS = 52 + EMG_ACC_GYRO_148HZ_16G_250DPS = 53 + EMG_ACC_GYRO_148HZ_2G_250DPS_2 = 54 + EMG_ACC_GYRO_148HZ_4G_250DPS = 55 + EMG_ACC_GYRO_148HZ_8G_250DPS = 56 + EMG_ACC_GYRO_148HZ_16G_250DPS = 57 + EMG_ACC_GYRO_148HZ_2G_500DPS = 58 + EMG_ACC_GYRO_148HZ_4G_500DPS = 59 + EMG_ACC_GYRO_148HZ_8G_500DPS = 60 + EMG_ACC_GYRO_148HZ_16G_500DPS = 61 + EMG_ORIENTATION_74HZ_16BITS = 66 + EMG_ORIENTATION_74HZ_32BITS = 67 + EMG_RMS_ACC_GYRO_296HZ_2G_250DPS = 68 + EMG_RMS_ACC_GYRO_296HZ_4G_250DPS = 69 + EMG_RMS_ACC_GYRO_296HZ_8G_250DPS = 70 + EMG_RMS_ACC_GYRO_296HZ_16G_250DPS = 71 + EMG_RMS_ACC_GYRO_296HZ_2G_500DPS = 72 + EMG_RMS_ACC_GYRO_296HZ_4G_500DPS = 73 + EMG_RMS_ACC_GYRO_296HZ_8G_500DPS = 74 + EMG_RMS_ACC_GYRO_296HZ_16G_500DPS = 75 + EMG_RMS_ACC_GYRO_296HZ_2G_1000DPS = 76 + EMG_RMS_ACC_GYRO_296HZ_4G_1000DPS = 77 + EMG_RMS_ACC_GYRO_296HZ_8G_1000DPS = 78 + EMG_RMS_ACC_GYRO_296HZ_16G_1000DPS = 79 + EMG_RMS_ACC_GYRO_296HZ_2G_2000DPS = 80 + EMG_RMS_ACC_GYRO_296HZ_4G_2000DPS = 81 + EMG_RMS_ACC_GYRO_296HZ_8G_2000DPS = 82 + EMG_RMS_ACC_GYRO_296HZ_16G_2000DPS = 83 + EMG_4370HZ = 84 + EMG_ACC_GYRO_518HZ_2G_250DPS = 86 + EMG_ACC_GYRO_518HZ_4G_250DPS = 87 + EMG_ACC_GYRO_518HZ_8G_250DPS = 88 + EMG_ACC_GYRO_518HZ_16G_250DPS = 89 + EMG_ACC_GYRO_518HZ_2G_500DPS = 90 + EMG_ACC_GYRO_518HZ_4G_500DPS = 91 + EMG_ACC_GYRO_518HZ_8G_500DPS = 92 + EMG_ACC_GYRO_518HZ_16G_500DPS = 93 + EMG_ACC_GYRO_518HZ_2G_1000DPS = 94 + EMG_ACC_GYRO_518HZ_4G_1000DPS = 95 + EMG_ACC_GYRO_518HZ_8G_1000DPS = 96 + EMG_ACC_GYRO_518HZ_16G_1000DPS = 97 + EMG_ACC_GYRO_518HZ_2G_2000DPS = 98 + EMG_ACC_GYRO_518HZ_4G_2000DPS = 99 + EMG_ACC_GYRO_518HZ_8G_2000DPS = 100 + EMG_ACC_GYRO_518HZ_16G_2000DPS = 101 + EMG_ACC_GYRO_963HZ_2G_250DPS = 102 + EMG_ACC_GYRO_963HZ_4G_250DPS = 103 + EMG_ACC_GYRO_963HZ_8G_250DPS = 104 + EMG_ACC_GYRO_963HZ_16G_250DPS = 105 + EMG_ACC_GYRO_963HZ_2G_500DPS = 106 + EMG_ACC_GYRO_963HZ_4G_500DPS = 107 + EMG_ACC_GYRO_963HZ_8G_500DPS = 108 + EMG_ACC_GYRO_963HZ_16G_500DPS = 109 + EMG_ACC_GYRO_963HZ_2G_1000DPS = 110 + EMG_ACC_GYRO_963HZ_4G_1000DPS = 111 + EMG_ACC_GYRO_963HZ_8G_1000DPS = 112 + EMG_ACC_GYRO_963HZ_16G_1000DPS = 113 + EMG_ACC_GYRO_963HZ_2G_2000DPS = 114 + EMG_ACC_GYRO_963HZ_4G_2000DPS = 115 + EMG_ACC_GYRO_963HZ_8G_2000DPS = 116 + EMG_ACC_GYRO_963HZ_16G_2000DPS = 117 + EMG_ACC_GYRO_741HZ_2G_250DPS = 118 + EMG_ACC_GYRO_741HZ_4G_250DPS = 119 + EMG_ACC_GYRO_741HZ_8G_250DPS = 120 + EMG_ACC_GYRO_741HZ_16G_250DPS = 121 + EMG_ACC_GYRO_741HZ_2G_500DPS = 122 + EMG_ACC_GYRO_741HZ_4G_500DPS = 123 + EMG_ACC_GYRO_741HZ_8G_500DPS = 124 + EMG_ACC_GYRO_741HZ_16G_500DPS = 125 + EMG_ACC_GYRO_741HZ_2G_1000DPS = 126 + EMG_ACC_GYRO_741HZ_4G_1000DPS = 127 + EMG_ACC_GYRO_741HZ_8G_1000DPS = 128 + EMG_ACC_GYRO_741HZ_16G_1000DPS = 129 + EMG_ACC_GYRO_741HZ_2G_2000DPS = 130 + EMG_ACC_GYRO_741HZ_4G_2000DPS = 131 + EMG_ACC_GYRO_741HZ_8G_2000DPS = 132 + EMG_ACC_GYRO_741HZ_16G_2000DPS = 133 + EMG_ORIENTATION_2370HZ_222HZ_32BITS = 134 + EMG_ACC_GYRO_74HZ_2G_250DPS_4000HZ = 153 + EMG_ACC_GYRO_74HZ_4G_250DPS_4000HZ = 154 + EMG_ACC_GYRO_74HZ_8G_250DPS_4000HZ = 155 + EMG_ACC_GYRO_74HZ_16G_250DPS_4000HZ = 156 + EMG_ACC_GYRO_74HZ_2G_500DPS_4000HZ = 157 + EMG_ACC_GYRO_74HZ_4G_500DPS_4000HZ = 158 + EMG_ACC_GYRO_74HZ_8G_500DPS_4000HZ = 159 + EMG_ACC_GYRO_74HZ_16G_500DPS_4000HZ = 160 + EMG_ACC_GYRO_74HZ_2G_1000DPS_4000HZ = 161 + EMG_ACC_GYRO_74HZ_4G_1000DPS_4000HZ = 162 + EMG_ACC_GYRO_74HZ_8G_1000DPS_4000HZ = 163 + EMG_ACC_GYRO_74HZ_16G_1000DPS_4000HZ = 164 + EMG_ACC_GYRO_74HZ_2G_2000DPS_4000HZ = 165 + EMG_ACC_GYRO_74HZ_4G_2000DPS_4000HZ = 166 + EMG_ACC_GYRO_74HZ_8G_2000DPS_4000HZ = 167 + EMG_ACC_GYRO_74HZ_16G_2000DPS_4000HZ = 168 + EMG_ORIENTATION_74HZ_16BITS_4000HZ = 169 + EMG_ORIENTATION_74HZ_32BITS_3740HZ = 170 + + @classmethod + def has_value(cls, value): + return any(value == item.value for item in cls) diff --git a/biosiglive/interfaces/trigno_sdk/enums.py b/biosiglive/interfaces/trigno_sdk/enums.py new file mode 100644 index 0000000..0f7d0e5 --- /dev/null +++ b/biosiglive/interfaces/trigno_sdk/enums.py @@ -0,0 +1,36 @@ +from enum import Enum + + +class AvantiAdapter(Enum): + """ + Return mode of sensor + """ + + Gogniometer = "23" + NONE = "O" + + +class AvantiSensor: + def __init__(self, adapter: AvantiAdapter = AvantiAdapter.NONE): + self.emg_port = 50043 + self.aux_port = 50044 + self.type = adapter.value + + +class LegacySensor: + def __init__(self): + self.emg_port = 50041 + self.aux_port = 50042 + self.type = "A" + + +class SensorType(Enum): + AVANTI = "Avanti" + LEGACY = "Legacy" + # ALL = "All" + # Avanti_all = "Avanti_all" + # Avanti_emg = "Avanti_emg" + # Avanti_aux = "Avanti_aux" + # Legacy_all = "Legacy_all" + # Legacy_emg = "Legacy_emg" + # Legacy_aux = "Legacy_aux" diff --git a/biosiglive/interfaces/trigno_sdk/gognio_mode_enum.py b/biosiglive/interfaces/trigno_sdk/gognio_mode_enum.py new file mode 100644 index 0000000..bbba37d --- /dev/null +++ b/biosiglive/interfaces/trigno_sdk/gognio_mode_enum.py @@ -0,0 +1,47 @@ +from enum import Enum + + +class GoniometerMode(Enum): + MODE_362 = 362 # SIG x2 @296Hz, ACC 2g, GYRO 250dps + MODE_363 = 363 # SIG x2 @296Hz, ACC 4g, GYRO 250dps + MODE_364 = 364 # SIG x2 @296Hz, ACC 8g, GYRO 250dps + MODE_365 = 365 # SIG x2 @296Hz, ACC 16g, GYRO 250dps + MODE_366 = 366 # SIG x2 @296Hz, ACC 2g, GYRO 500dps + MODE_367 = 367 # SIG x2 @296Hz, ACC 4g, GYRO 500dps + MODE_368 = 368 # SIG x2 @296Hz, ACC 8g, GYRO 500dps + MODE_369 = 369 # SIG x2 @296Hz, ACC 16g, GYRO 500dps + MODE_370 = 370 # SIG x2 @296Hz, ACC 2g, GYRO 1000dps + MODE_371 = 371 # SIG x2 @296Hz, ACC 4g, GYRO 1000dps + MODE_372 = 372 # SIG x2 @296Hz, ACC 8g, GYRO 1000dps + MODE_373 = 373 # SIG x2 @296Hz, ACC 16g, GYRO 1000dps + MODE_374 = 374 # SIG x2 @296Hz, ACC 2g, GYRO 2000dps + MODE_375 = 375 # SIG x2 @296Hz, ACC 4g, GYRO 2000dps + MODE_376 = 376 # SIG x2 @296Hz, ACC 8g, GYRO 2000dps + MODE_377 = 377 # SIG x2 @296Hz, ACC 16g, GYRO 2000dps + MODE_378 = 378 # SIG x2 @370Hz, OR 32-bit @74Hz + MODE_026 = 26 # 1 HF Chan @1926Hz, 1 LF Chan @148Hz + MODE_244 = 244 # SIG x2 @519Hz + + def description(self): + descriptions = { + GoniometerMode.MODE_362: "SIG x2 @296Hz, ACC 2g, GYRO 250dps", + GoniometerMode.MODE_363: "SIG x2 @296Hz, ACC 4g, GYRO 250dps", + GoniometerMode.MODE_364: "SIG x2 @296Hz, ACC 8g, GYRO 250dps", + GoniometerMode.MODE_365: "SIG x2 @296Hz, ACC 16g, GYRO 250dps", + GoniometerMode.MODE_366: "SIG x2 @296Hz, ACC 2g, GYRO 500dps", + GoniometerMode.MODE_367: "SIG x2 @296Hz, ACC 4g, GYRO 500dps", + GoniometerMode.MODE_368: "SIG x2 @296Hz, ACC 8g, GYRO 500dps", + GoniometerMode.MODE_369: "SIG x2 @296Hz, ACC 16g, GYRO 500dps", + GoniometerMode.MODE_370: "SIG x2 @296Hz, ACC 2g, GYRO 1000dps", + GoniometerMode.MODE_371: "SIG x2 @296Hz, ACC 4g, GYRO 1000dps", + GoniometerMode.MODE_372: "SIG x2 @296Hz, ACC 8g, GYRO 1000dps", + GoniometerMode.MODE_373: "SIG x2 @296Hz, ACC 16g, GYRO 1000dps", + GoniometerMode.MODE_374: "SIG x2 @296Hz, ACC 2g, GYRO 2000dps", + GoniometerMode.MODE_375: "SIG x2 @296Hz, ACC 4g, GYRO 2000dps", + GoniometerMode.MODE_376: "SIG x2 @296Hz, ACC 8g, GYRO 2000dps", + GoniometerMode.MODE_377: "SIG x2 @296Hz, ACC 16g, GYRO 2000dps", + GoniometerMode.MODE_378: "SIG x2 @370Hz, OR 32-bit @74Hz", + GoniometerMode.MODE_026: "1 HF Chan @1926Hz, 1 LF Chan @148Hz", + GoniometerMode.MODE_244: "SIG x2 @519Hz", + } + return descriptions.get(self, "Unknown Mode") diff --git a/biosiglive/interfaces/trigno_sdk/sdk_client.py b/biosiglive/interfaces/trigno_sdk/sdk_client.py new file mode 100644 index 0000000..8c91f0b --- /dev/null +++ b/biosiglive/interfaces/trigno_sdk/sdk_client.py @@ -0,0 +1,443 @@ +import socket +import struct +import threading +from queue import Queue +import time + +import numpy as np + +from .enums import AvantiSensor, LegacySensor, Enum +from .sensor import Sensor, Type + + +BYTES_PER_CHANNEL = 4 +CMD_TERM = "\r\n\r\n" +EMG_SAMPLE_RATE = 2000 +AUX_SAMPLE_RATE = 148.148 +SAMPLE_INTERVAL = 0.0135 + + +class TrignoSDKClient: + def __init__( + self, + host="127.0.0.1", + cmd_port=50040, + timeout=2.0, + fast_mode=False, + buffer_size=1000, + stream_rate=74.074074, + init_sensors=True, + ): + self.buffer_size = buffer_size + self.targeted_rate = stream_rate + self.host = host + self.cmd_port = cmd_port + self.timeout = timeout + self.is_connected = False + self._comm_socket = None + self.fast_mode = fast_mode + self.avanti_emg_socket = None + self.avanti_aux_socket = None + self.legacy_emg_socket = None + self.legacy_aux_socket = None + self.time_counter = time.perf_counter + self.all_socket = {} + self.all_events = {} + self._last_data = { + "avanti_emg": None, + "avanti_aux": None, + "legacy_emg": None, + "legacy_aux": None, + } + + self.avanti_emg_queue = Queue() + self.avanti_aux_queue = Queue() + self.legacy_emg_queue = Queue() + self.legacy_aux_queue = Queue() + self.all_queue = { + "avanti_emg": self.avanti_emg_queue, + "avanti_aux": self.avanti_aux_queue, + "legacy_emg": self.legacy_emg_queue, + "legacy_aux": self.legacy_aux_queue, + } + + if self.fast_mode: + print("Warning: Fast mode enabled. Responses will not be waited for.") + + self.connect(init_sensors=init_sensors) + + def connect(self, init_sensors=True): + """Establish connection to Trigno SDK command port.""" + self._comm_socket = socket.create_connection((self.host, self.cmd_port), self.timeout) + self.is_connected = True + self.send_command("BACKWARDS COMPATIBILITY OFF") + if init_sensors: + self.initialize_sensors() + + try: + _ = self._comm_socket.recv(1024) + except socket.timeout: + pass + + def reconnect_device(self): + for _, item in self.all_socket.items(): + if item is not None: + item.close() + for _, event in self.all_events.items(): + if event is not None: + event = None + self.initialize_sensors() + + def initiate_data_connection(self): + for key, item in self._threads_to_run.items(): + self.all_events[key] = item if not item else threading.Event() + self.all_socket[key] = item if not item else self._connect_to_socket(self._port_from_name(key)) + + def _port_from_name(self, name): + if "emg" in name: + return AvantiSensor().emg_port if "avanti" in name.lower() else LegacySensor().emg_port + elif "aux" in name: + return AvantiSensor().aux_port if "avanti" in name.lower() else LegacySensor().aux_port + else: + raise ValueError("Type not recognized.") + + def _check_data_to_stream(self, data_to_stream): + if "emg" not in data_to_stream and "delsys_gogniometer" not in data_to_stream: + self._threads_to_run["avanti_emg"] = False + self._threads_to_run["legacy_emg"] = False + if "aux" not in data_to_stream and "delsys_gogniometer" not in data_to_stream: + self._threads_to_run["avanti_aux"] = False + self._threads_to_run["legacy_aux"] = False + + def initialize_sensors(self, data_to_stream=None): + """Initialize all sensors.""" + self.sensors = [Sensor(i, self, self.buffer_size) for i in range(1, 17)] + self._get_which_thread_to_run() + if data_to_stream is not None: + self._check_data_to_stream(data_to_stream) + self.initiate_data_connection() + + def _get_which_thread_to_run(self): + all_emg = [] + all_aux = [] + for sensor in self.sensors: + if not sensor.is_paired: + continue + emg, aux = sensor.emg_aux_type() + all_emg.append(emg) + all_aux.append(aux) + + self._threads_to_run = { + "avanti_emg": True in ["Avanti" in _.name for _ in all_emg if _ is not None], + "avanti_aux": True in ["Avanti" in _.name for _ in all_aux if _ is not None], + "legacy_emg": True in ["Legacy" in _.name for _ in all_emg if _ is not None], + "legacy_aux": True in ["Legacy" in _.name for _ in all_aux if _ is not None], + } + + def _connect_to_socket(self, port): + _data_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + _data_socket.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1) + _data_socket.connect((self.host, port)) + _data_socket.setblocking(False) + return _data_socket + + def _bytes_per_sample(self, n_channels): + """Return the number of bytes per sample for a given number of channels.""" + return n_channels * BYTES_PER_CHANNEL + + def buffer_size(self, n_channels, n_samples): + """Return the size of the buffer required to store a given number of samples for a given number of channels.""" + return self._bytes_per_sample(n_channels) * n_samples + + def disconnect(self): + """Close the socket connection.""" + if self._comm_socket: + self._comm_socket.close() + self._comm_socket = None + + def send_command(self, command: str) -> str: + """ + Send a command or query to the Trigno system and return the response as a string. + Command strings must already include any needed arguments. + """ + if self._comm_socket is None: + raise RuntimeError("Not connected. Call connect() first.") + + full_command = f"{command}\r\n\r\n" + self._comm_socket.sendall(full_command.encode("ascii")) + + # If in fast mode, return immediately without waiting for a response + if self.fast_mode: + return "" + + # Give the server time to respond + while True: + time.sleep(0.1) + try: + response = self._comm_socket.recv(1024) + return response.decode("ascii").strip() + except socket.timeout: + None + + def read(self, connection, buffer_size, n_channels, sample_interval, data_matrix=None): + l = -1 + packet = bytes() + flush_packet = bytes() + time_counter = self.time_counter() + try: + while l != 0: + flush_packet = connection.recv(4096) + packet += flush_packet + l = len(flush_packet) + except BlockingIOError: + pass + + if data_matrix is None: + data_matrix = np.zeros((int(buffer_size // 4 / n_channels), n_channels)) + data = np.asarray(struct.unpack("<" + "f" * (len(packet) // 4), packet)) + nb_flushed_data = int(buffer_size // 4 / n_channels) - len(data) // n_channels + if nb_flushed_data > 0: + time_counter += nb_flushed_data * sample_interval + data_matrix[: len(data) // n_channels, :] = data.reshape((-1, n_channels))[ + -int(buffer_size // 4 / n_channels) :, : + ] + + return data_matrix.T, time_counter + + def start_streaming(self): + # self.streaming_rate = self._compute_rate_from_sensors(rate) + start_response = self.send_command("START") + is_started = start_response == "OK" + if not is_started and not self.fast_mode: + raise RuntimeError("Streaming not started.") + self._launch_threads() + + def _compute_rate_from_sensors(self, initial_rate): + # get minimum max samples + closest_rate = initial_rate + sample_numbers = min(self.get_max_aux_samples(), self.get_max_emg_samples) + n_sample_init = sample_numbers * (1 / initial_rate) / SAMPLE_INTERVAL + if n_sample_init % 1 != 0: + closest_rate = 1 / (int(n_sample_init) * SAMPLE_INTERVAL / sample_numbers) + print( + f"WARNING: You requested an unvalid sampling rate ({initial_rate} Hz)." + f"The rate was set to the closest valid rate: {closest_rate:.3f} Hz." + ) + + return closest_rate + + def buffer_size_for_type(self, name, n_samples=None): + if "emg" in name: + n_samples = self.get_max_emg_samples() if not n_samples else n_samples + n_channel = 16 + buffer_size = n_channel * n_samples * BYTES_PER_CHANNEL + elif "aux" in name: + n_samples = self.get_max_aux_samples() if not n_samples else n_samples + n_channel = 144 if "avanti" in name.lower() else 48 + buffer_size = n_channel * n_samples * BYTES_PER_CHANNEL + else: + raise RuntimeError("Invalid sensor type.") + return buffer_size, n_channel, n_samples + + @staticmethod + def flush_socket(socket): + try: + while True: + socket.recv(4096) + except BlockingIOError: + return + + def _launch_one_thread(self, socket_tmp, name, data_queue, event): + buffer_size, n_channels, n_samples = self.buffer_size_for_type(name) + data_matrix = np.zeros((int(buffer_size // 4 / n_channels), n_channels)) + channel_native_streaming_rate = ( + self.get_emg_streaming_rate() if "emg" in name else self.get_aux_streaming_rate() + ) + sample_interval = 1 / channel_native_streaming_rate + streaming_interval = self._get_interval_from_channel(name) + self.flush_socket(socket_tmp) + + def _read_function(): + data, timestamp = self.read(socket_tmp, buffer_size, n_channels, sample_interval, data_matrix) + try: + data_queue.get_nowait() + except: + pass + data_queue.put_nowait((data[np.unique(data.nonzero()[0])], timestamp)) + + def _thread_func(): + while True: + # tic = self.time_counter() + _read_function() + # toc = self.time_counter() + # if toc - tic < streaming_interval: + # time.sleep(streaming_interval - (toc - tic)) + # else: + # print(f"Warning: Streaming interval exceeded. Time taken to read data: {toc - tic:.3f} s, target time: {streaming_interval:.3f} s.") + + thread = threading.Thread(target=_thread_func, name=name) + thread.start() + + def _get_interval_from_channel(self, name): + max_samples = self.get_max_emg_samples() if "emg" in name else self.get_max_aux_samples() + sample = int(max_samples * (1 / self.targeted_rate) / SAMPLE_INTERVAL) + return sample * SAMPLE_INTERVAL / max_samples + + def _launch_threads(self): + all_soc, all_q, all_ev = self.all_socket, self.all_queue, self.all_events + self._all_threads = [ + self._launch_one_thread(all_soc[n], n, all_q[n], all_ev[n]) + for n in self._threads_to_run.keys() + if self._threads_to_run[n] + ] + + # def _main_thread_func(): + # while True: + # for name in self.all_socket.keys(): + # self.all_events[name].wait() + # self.all_events[name].clear() + # self._set_all_data() + + # main_thread = threading.Thread(target=_main_thread_func, name='main') + # main_thread.start() + + def stop_streaming(self): + return self.send_command("STOP") + + def disconnect(self): + self.stop_streaming() + _ = [socket.close() for socket in self.all_socket.values()] + self._comm_socket.close() + + def get_emg_streaming_rate(self): + return float(self.send_command("MAX SAMPLES EMG")) / 0.0135 + + def get_aux_streaming_rate(self): + return float(self.send_command("MAX SAMPLES AUX")) / 0.0135 + + def get_max_emg_samples(self): + return int(self.send_command("MAX SAMPLES EMG")) + + def get_max_aux_samples(self): + return int(self.send_command("MAX SAMPLES AUX")) + + def get_aux_streaming_rate(self): + return int(self.send_command("MAX SAMPLES AUX")) / 0.0135 + + def get_trigger_state(self): + return self.send_command("TRIGGER?") + + def set_trigger(self, which="START", state="ON"): + return self.send_command(f"TRIGGER {which} {state}") + + def get_backwards_compatibility(self): + return self.send_command("BACKWARDS COMPATIBILITY?") + + def set_backwards_compatibility(self, state="ON"): + return self.send_command(f"BACKWARDS COMPATIBILITY {state}") + + def get_upsampling(self): + return self.send_command("UPSAMPLING?") + + def set_upsampling(self, state="ON"): + return self.send_command(f"UPSAMPLE {state}") + + def get_sensor_info(self, n, info="TYPE"): + return self.send_command(f"SENSOR {n} {info}?") + + def get_sensor_emgchannel(self, n): + if not self.is_sensor_paired(n): + return 0 + return int(self.send_command(f"SENSOR {n} EMGCHANNELCOUNT?")) + + def get_sensor_auxchannel(self, n): + if not self.is_sensor_paired(n): + return 0 + return int(self.send_command(f"SENSOR {n} AUXCHANNELCOUNT?")) + + def is_sensor_paired(self, n): + return self.send_command(f"SENSOR {n} PAIRED?") == "YES" + + def get_sensor_idx(self, n): + return int(self.send_command(f"SENSOR {n} STARTINDEX?")) + + def get_list_sensors_and_idx(self): + sensors = [] + for i in range(1, 16): + if self.is_sensor_paired(i): + sensors.append([self.get_sensor_info(i, "TYPE"), self.get_sensor_idx(i)]) + else: + sensors.append([None, None]) + return sensors + + def pair_sensor(self, n): + return self.send_command(f"SENSOR {n} PAIR") + + def set_sensor_mode(self, n, mode): + return self.send_command(f"SENSOR {n} SETMODE {mode}") + + def get_endianness(self): + return self.send_command("ENDIANNESS?") + + def set_endianness(self, mode="LITTLE"): + return self.send_command(f"ENDIAN {mode}") + + def get_base_serial(self): + return self.send_command("BASE SERIAL?") + + def get_base_firmware(self): + return self.send_command("BASE FIRMWARE?") + + def get_number_emgchannel(self): + nb_channel = 0 + for i in range(1, 16): + nb_channel += self.get_sensor_emgchannel(i) + return nb_channel + + def get_number_auxchannel(self): + nb_channel = 0 + for i in range(1, 16): + nb_channel += self.get_sensor_auxchannel(i) + return nb_channel + + def _set_avanti_emg_data(self): + try: + data = self.avanti_emg_queue.get_nowait() + for sensor in self.sensors: + if sensor.type == Type.Avanti or sensor.type == Type.AvantiGogniometer: + sensor.update_emg_buffer(data[0][sensor.emg_range[0] : sensor.emg_range[1], :], data[1]) + except: + return + + def _set_avanti_aux_data(self): + try: + data = self.avanti_aux_queue.get_nowait() + for sensor in self.sensors: + if sensor.type == Type.Avanti or sensor.type == Type.AvantiGogniometer: + sensor.update_aux_buffer(data[0][sensor.aux_range[0] : sensor.aux_range[1], :], data[1]) + except: + return + + def _set_legacy_emg_data(self): + try: + data = self.legacy_emg_queue.get_nowait() + for sensor in self.sensors: + if sensor.type == Type.Legacy: + sensor.update_emg_buffer(data[0][sensor.emg_range[0] : sensor.emg_range[1], :], data[1]) + except: + return + + def _set_legacy_aux_data(self): + try: + data = self.legacy_aux_queue.get_nowait() + for sensor in self.sensors: + if sensor.type == Type.Legacy: + sensor.update_aux_buffer(data[0][sensor.aux_range[0] : sensor.aux_range[1], :], data[1]) + except: + return + + def _set_all_data(self): + self._set_avanti_emg_data(), + self._set_avanti_aux_data(), + self._set_legacy_emg_data(), + self._set_legacy_aux_data(), diff --git a/biosiglive/interfaces/trigno_sdk/sensor.py b/biosiglive/interfaces/trigno_sdk/sensor.py new file mode 100644 index 0000000..baa70ec --- /dev/null +++ b/biosiglive/interfaces/trigno_sdk/sensor.py @@ -0,0 +1,155 @@ +from enum import Enum +import numpy as np + + +class Type(Enum): + Avanti = "O" + Legacy = "A" + AvantiGogniometer = "23" + + +class Sensor: + def __init__(self, index, trigno_box=None, buff_size=100): + self.buff_size = buff_size + self.name = f"sensor {index}" + self.index = index + self.type = None + self.mode = None + self.units = None + self.type = None + self.range = None + self.nb_emg_channels = None + self.nb_aux_channels = None + self.max_emg_samples = None + self.max_aux_samples = None + self.emg_rate = None + self.aux_rate = None + self._index_emg = 0 + self._index_aux = 0 + self._emg_frame_numbers = [] + self._aux_frame_numbers = [] + self.sensor_start_idx = 0 + + self.emg_buffer = None + self.aux_buffer = None + self.trigno_box = trigno_box + + if trigno_box is not None: + self.initialize() + + def emg_aux_type(self): + if self.is_paired: + emg = self.type if self.nb_emg_channels > 0 else None + aux = self.type if self.nb_aux_channels > 0 else None + return emg, aux + else: + return None, None + + def initialize(self): + """ + Initialize the sensor + :param trigno_box: TrignoBox object + :return: None + """ + if not self.is_sensor_paired(): + self.is_paired = False + return + + self.is_paired = True + self.mode = self.get_sensor_mode() + # self.units = self.get_sensor_units() + self.type = self.get_sensor_type() + self.nb_emg_channels = self.get_sensor_emgchannel() + self.nb_aux_channels = self.get_sensor_auxchannel() + self.aux_rate = self.get_aux_streaming_rate() + self.emg_rate = self.get_emg_streaming_rate() + self.sensor_start_idx = self.get_sensor_idx() + self.emg_range = (self.sensor_start_idx, self.sensor_start_idx + self.nb_emg_channels) + self.aux_range = (self.sensor_start_idx * 9, self.sensor_start_idx * 9 + self.nb_aux_channels) + + # self.emg_buffer = np.empty((self.nb_emg_channels, self.max_emg_samples, self.buff_size)) + # self.aux_buffer = np.empty((self.nb_aux_channels, self.max_aux_samples, self.buff_size)) + + @property + def last_emg_chunck(self): + return self.emg_buffer[..., self._index_emg] + + @property + def last_aux_chunck(self): + return self.aux_buffer[..., self._index_aux] + + def update_emg_buffer(self, emg_data, n_chunck=None): + if not self.is_paired: + return + self.extend_frame_numbers(self._emg_frame_numbers, n_chunck, self.max_emg_samples) + self.emg_buffer[..., self._index_emg] = emg_data + self._index_emg = (self._index_emg + 1) % self.buff_size + + def get_emg_streaming_rate(self): + if int(self.trigno_box.send_command(f"SENSOR {self.index} EMGCHANNELCOUNT?")) > 0: + return float(self.trigno_box.send_command(f"SENSOR {self.index} CHANNEL 1 RATE?")) + else: + return None + + def get_aux_streaming_rate(self): + if int(self.trigno_box.send_command(f"SENSOR {self.index} AUXCHANNELCOUNT?")) > 0: + return float(self.trigno_box.send_command(f"SENSOR {self.index} CHANNEL 2 RATE?")) + else: + return None + + def update_aux_buffer(self, aux_data, n_chunck=None): + if not self.is_paired: + return + self.extend_frame_numbers(self._aux_frame_numbers, n_chunck, self.max_aux_samples) + self.aux_buffer[..., self._index_aux] = aux_data + self._index_aux = (self._index_aux + 1) % self.buff_size + + @property + def emg_frame_numbers(self): + return self._emg_frame_numbers + + @property + def aux_frame_numbers(self): + return self._aux_frame_numbers + + def extend_frame_numbers(self, init_list, n_chunck, n_samples): + if n_chunck is None: + init_list.extend(list(range(len(init_list), len(init_list) + n_samples))) + else: + init_list.extend(list(range(n_chunck, n_chunck + n_samples))) + return init_list + + def get_emg_from_buffer(self): + return np.roll(self.emg_buffer, -self._index_emg, axis=-1).reshape((self.nb_emg_channels, -1)) + + def get_aux_from_buffer(self): + return np.roll(self.aux_buffer, -self._index_aux, axis=-1).reshape((self.nb_aux_channels, -1)) + + def get_sensor_info(self, info="TYPE"): + return self.trigno_box.send_command(f"SENSOR {self.index} {info}?") + + def get_sensor_type(self): + sensor_type = self.trigno_box.send_command(f"SENSOR {self.index} TYPE?") + try: + return Type(sensor_type) + except ValueError: + raise ValueError(f"Sensor {self.index} has an unknown type: {sensor_type}") + + def get_sensor_mode(self): + return self.trigno_box.send_command(f"SENSOR {self.index} MODE?") + + def get_sensor_emgchannel(self): + if not self.is_sensor_paired(): + return 0 + return int(self.trigno_box.send_command(f"SENSOR {self.index} EMGCHANNELCOUNT?")) + + def get_sensor_auxchannel(self): + if not self.is_sensor_paired(): + return 0 + return int(self.trigno_box.send_command(f"SENSOR {self.index} AUXCHANNELCOUNT?")) + + def is_sensor_paired(self): + return self.trigno_box.send_command(f"SENSOR {self.index} PAIRED?") == "YES" + + def get_sensor_idx(self): + return int(self.trigno_box.send_command(f"SENSOR {self.index} STARTINDEX?")) diff --git a/biosiglive/interfaces/vicon_interface.py b/biosiglive/interfaces/vicon_interface.py index 51c4386..51d9543 100644 --- a/biosiglive/interfaces/vicon_interface.py +++ b/biosiglive/interfaces/vicon_interface.py @@ -1,6 +1,7 @@ """ This file contains a wrapper for the python Vicon SDK. """ + from .param import * from typing import Union from .generic_interface import GenericInterface @@ -63,20 +64,22 @@ def _init_client(self): self.vicon_client.EnableMarkerData() self.vicon_client.EnableUnlabeledMarkerData() self.get_frame() - if self.system_rate != self.vicon_client.GetFrameRate(): + if self.system_rate is None: + self.system_rate = self.vicon_client.GetFrameRate() + elif self.system_rate != self.vicon_client.GetFrameRate(): raise ValueError( f"Vicon system rate ({self.vicon_client.GetFrameRate()}) does not match the system rate " f"({self.system_rate})." ) + # TODO: Remove rate and get it from the Vicon system. def add_device( self, nb_channels: int, device_type: Union[DeviceType, str] = DeviceType.Emg, data_buffer_size: int = None, name: str = None, - rate: float = 2000, - device_range: tuple = None, + rate: float = None, processing_method: Union[RealTimeProcessingMethod, OfflineProcessingMethod] = None, **process_kwargs, ): @@ -102,12 +105,23 @@ def add_device( **process_kwargs Keyword arguments for the processing method. """ - device_tmp = self._add_device( - nb_channels, device_type, name, rate, device_range, processing_method, **process_kwargs - ) + device_tmp = self._add_device(nb_channels, device_type, name, rate, processing_method, **process_kwargs) device_tmp.interface = self.interface_type if self.vicon_client: - device_tmp.infos = self.vicon_client.GetDeviceOutputDetails(name) + for i in range(10): + try: + device_tmp.infos = self.vicon_client.GetDeviceOutputDetails(name) + except: + self.get_frame() + pass + if device_tmp.infos is None: + raise RuntimeError( + f"Device {name} not found on Vicon. Available devices are {self.vicon_client.GetDeviceNames()}." + ) + # self.get_frame() + # device_tmp.infos = self.vicon_client.GetDeviceOutputDetails(name) + + # raise RuntimeError(f"Device {name} not found on Vicon. Available devices are {self.vicon_client.GetDeviceNames()}.") else: device_tmp.infos = None device_tmp.data_windows = data_buffer_size @@ -232,8 +246,15 @@ def get_device_data( all_device_data.append(device_data) if len(all_device_data) == 1: return all_device_data[0] + if len(all_device_data) == 0: + return None return all_device_data + def get_timestamp(self, get_frame=False): + if get_frame: + self.get_frame() + return self.vicon_client.GetTimecode() + def get_marker_set_data( self, subject_name: Union[str, list] = None, marker_names: Union[str, list] = None, get_frame: bool = True ): @@ -254,6 +275,9 @@ def get_marker_set_data( markers_data: list All asked markers data. """ + if get_frame: + self.get_frame() + if len(self.marker_sets) == 0: raise ValueError("No marker set has been added to the Vicon system.") if not self.is_initialized: diff --git a/biosiglive/processing/data_processing.py b/biosiglive/processing/data_processing.py index c4050fa..24e3140 100644 --- a/biosiglive/processing/data_processing.py +++ b/biosiglive/processing/data_processing.py @@ -146,7 +146,7 @@ def _moving_average(data: np.ndarray, window: np.ndarray, empty_ma: np.ndarray) Moving average. """ for i in range(data.shape[0]): - empty_ma[i, :] = convolve(data[i, :], window, mode="same", method="fft") + empty_ma[i, :] = convolve(data[i, :], window, mode="same", method="auto") return empty_ma @staticmethod @@ -326,6 +326,7 @@ def process_emg( absolute_value=True, normalization=False, moving_average_window=200, + window_weights: list = None, **kwargs, ) -> np.ndarray: """ @@ -358,12 +359,13 @@ def process_emg( processed EMG data. """ + emg_data = emg_data.copy() self.update_signal_processing_parameters(**kwargs) tic = time.time() if low_pass_filter and moving_average: raise RuntimeError("Please choose between low-pass filter and moving average.") ma_win = moving_average_window - if ma_win > self.processing_window: + if ma_win > self.processing_window and moving_average: raise RuntimeError(f"Moving average windows ({ma_win}) higher than emg windows ({self.processing_window}).") emg_sample = emg_data.shape[1] if emg_sample == 0: @@ -421,7 +423,7 @@ def process_emg( average = np.median(emg_proc_tmp[:, -ma_win:], axis=1).reshape(-1, 1) self.processed_data_buffer = np.append(self.processed_data_buffer[:, 1:], average / quot, axis=1) self.process_time.append(time.time() - tic) - return self.processed_data_buffer + return self.processed_data_buffer.copy() def process_imu( self, diff --git a/biosiglive/processing/msk_functions.py b/biosiglive/processing/msk_functions.py index c4c5095..b3ee7af 100644 --- a/biosiglive/processing/msk_functions.py +++ b/biosiglive/processing/msk_functions.py @@ -1,6 +1,7 @@ """ This file contains biorbd specific functions for musculoskeletal analysis such as inverse or direct kinematics. """ + try: import biorbd @@ -9,22 +10,48 @@ biordb_package = False import numpy as np from ..enums import InverseKinematicsMethods +from .data_processing import RealTimeProcessing from typing import Union import time +try: + import casadi as ca + from acados_template import AcadosOcp, AcadosModel, AcadosOcpSolver + from .msk_utils import ( + _init_acados, + _update_solver, + _init_casadi_function, + _create_new_model, + _compute_forces, + _compute_inverse_dynamics, + _express_in_new_coordinate, + ExternalLoads, + ) +except ModuleNotFoundError: + pass + class MskFunctions: - def __init__(self, model: str, data_buffer_size: int = 1): + def __init__(self, model: str, data_buffer_size: int = 1, system_rate: int = 100): """ The MskFunctions contains all function for some musculoskeletal methods. Parameters ---------- - model : str + model : Union[str, biorbd.Model] Path to the biorbd model used to compute the kinematics. data_buffer_size: int The size of the buffer used to store the data. + system_rate: int + The working frequency of the input data (markers or joint kinematics). """ + self.weigh_list = None + self.model_all_dofs = None + self.ca_model = None + self.once_compile = None + self.ocp = None + self.mjt_funct = None + self.ocp_solver = None if not biordb_package: raise ModuleNotFoundError( "Biorbd is not installed." @@ -38,8 +65,22 @@ def __init__(self, model: str, data_buffer_size: int = 1): self.process_time = [] self.markers_buffer = [] self.kin_buffer = [] + self.dyn_buffer = [] + self.act_buffer = [] + self.id_state_buffer = [] + self.tau_buffer = [] + self.jrf_buffer = [] + self.res_tau_buffer = [] self.data_windows = data_buffer_size + self.system_rate = system_rate self.kalman = None + self._rt_process_methods = None + self._state_idx_to_process = None + self.q_mapping, self.ordered_seg, self.ordered_idx = None, None, None + + def clean_all_buffers(self): + for key in self.__dict__.keys(): + self.__dict__[key] = [] if "buffer" in key else self.__dict__[key] def compute_inverse_kinematics( self, @@ -48,6 +89,10 @@ def compute_inverse_kinematics( kalman_freq: Union[int, float] = 100, kalman: callable = None, custom_function: callable = None, + initial_guess: Union[np.ndarray, list] = None, + qdot_from_finite_difference: bool = False, + noise_factor=1e-10, + error_factor=1e-5, **kwargs, ) -> tuple: """ @@ -65,6 +110,10 @@ def compute_inverse_kinematics( The method to use to compute the inverse kinematics. custom_function : callable Custom function to use. + initial_guess: Union[np.ndarray, list] + Initial generalized coordinate, velocity and acceleration for the kalman filter + qdot_from_finite_difference: bool + If true the velocity will be computed using a finite difference method. Returns ------- @@ -77,50 +126,74 @@ def compute_inverse_kinematics( method = InverseKinematicsMethods(method) else: raise ValueError(f"Method {method} is not supported") - + if qdot_from_finite_difference and markers.shape[2] < 2: + raise ValueError("You must have at least two frames to compute the velocity using finite difference.") if method == InverseKinematicsMethods.BiorbdKalman: self.kalman = kalman if kalman else self.kalman if not kalman and not self.kalman: freq = kalman_freq # Hz - params = biorbd.KalmanParam(freq) + params = biorbd.KalmanParam(freq, noiseFactor=noise_factor, errorFactor=error_factor) self.kalman = biorbd.KalmanReconsMarkers(self.model, params) - markers_over_frames = [] + if initial_guess: + if isinstance(initial_guess, np.ndarray): + if initial_guess.shape[0] != 3: + raise RuntimeError("Initial guess must have dims : 3xNdofxNframes if give as an array.") + initial_guess = [initial_guess[0, ...], initial_guess[1, ...], initial_guess[2, ...]] + for i in initial_guess: + if len(i.shape) != 1: + raise RuntimeError("initial guess must be 1D array.") + if i.shape[0] != self.model.nbQ(): + raise RuntimeError("Inital guess msut have the same size than model DOfs.") + if len(initial_guess) != 3: + raise RuntimeError("Initial guess must be of len 3 (angle, velocity, acceleration).") + self.kalman.setInitState(initial_guess[0], initial_guess[1], initial_guess[2]) + q = biorbd.GeneralizedCoordinates(self.model) q_dot = biorbd.GeneralizedVelocity(self.model) qd_dot = biorbd.GeneralizedAcceleration(self.model) + q_recons = np.zeros((self.model.nbQ(), markers.shape[2])) + q_dot_recons = np.zeros((self.model.nbQ(), markers.shape[2])) + q_ddot_recons = np.zeros((self.model.nbQ(), markers.shape[2])) for i in range(markers.shape[2]): - markers_over_frames.append([biorbd.NodeSegment(m) for m in markers[:, :, i].T]) - - q_recons = np.zeros((self.model.nbQ(), len(markers_over_frames))) - q_dot_recons = np.zeros((self.model.nbQ(), len(markers_over_frames))) - - for i, targetMarkers in enumerate(markers_over_frames): - self.kalman.reconstructFrame(self.model, targetMarkers, q, q_dot, qd_dot) + self.kalman.reconstructFrame( + self.model, [biorbd.NodeSegment(m) for m in markers[:, :, i].T], q, q_dot, qd_dot + ) q_recons[:, i] = q.to_array() q_dot_recons[:, i] = q_dot.to_array() - + q_ddot_recons[:, i] = qd_dot.to_array() elif method == InverseKinematicsMethods.BiorbdLeastSquare: ik = biorbd.InverseKinematics(self.model, markers) ik.solve("only_lm") + # ik.solve("trf") + q_recons = ik.q q_dot_recons = np.array([0] * ik.nb_q)[:, np.newaxis] + q_ddot_recons = np.array([0] * ik.nb_q)[:, np.newaxis] elif method == InverseKinematicsMethods.Custom: if not custom_function: raise ValueError("No custom function provided.") q_recons = custom_function(markers, **kwargs) - q_dot_recons = np.zerros((q_recons.shape())) + q_dot_recons = np.zeros((q_recons.shape())) + q_ddot_recons = np.zeros((q_recons.shape())) else: raise ValueError(f"Method {method} is not supported") + + if qdot_from_finite_difference: + q_dot_recons = np.zeros_like(q_recons) + for i in range(1, q_recons.shape[1] - 1): + q_dot_recons[:, i] = (q_recons[:, i + 1] - q_recons[:, i - 1]) / (2 / self.system_rate) + if len(self.kin_buffer) == 0: - self.kin_buffer = [q_recons, q_dot_recons] + self.kin_buffer = [q_recons, q_dot_recons, q_ddot_recons] else: self.kin_buffer[0] = np.append(self.kin_buffer[0], q_recons, axis=1) self.kin_buffer[1] = np.append(self.kin_buffer[1], q_dot_recons, axis=1) + self.kin_buffer[2] = np.append(self.kin_buffer[2], q_ddot_recons, axis=1) for i in range(len(self.kin_buffer)): self.kin_buffer[i] = self.kin_buffer[i][:, -self.data_windows :] self.process_time.append(time.time() - tic) - return self.kin_buffer[0], self.kin_buffer[1] + return self.kin_buffer[0].copy(), self.kin_buffer[1].copy(), self.kin_buffer[2].copy() def compute_direct_kinematics(self, states: np.ndarray) -> np.ndarray: """ @@ -161,6 +234,486 @@ def compute_direct_kinematics(self, states: np.ndarray) -> np.ndarray: self.process_time.append(time.time() - tic) return self.markers_buffer + def compute_inverse_dynamics( + self, + joint_positions: np.ndarray = None, + joint_velocities: np.ndarray = None, + joint_accelerations: np.ndarray = None, + state_idx_to_process: list = (), + lowpass_frequency: Union[list, int] = None, + windows_length: Union[list, int] = None, + positions_from_inverse_kinematics: bool = False, + velocities_from_inverse_kinematics: bool = False, + accelerations_from_inverse_kinematics: bool = False, + external_load: any = None, + ) -> np.ndarray: + """ + Compute the inverse dynamics using the model kinematics and a biorbd model type. + + Parameters + ---------- + joint_positions : np.ndarray + The joint position for each model joints + joint_velocities : np.ndarray + The joint velocities for each model joints + joint_accelerations : np.ndarray + The joint accelerations for each model joints + state_idx_to_process: list + The list of the index of the data to apply a low pass filter on. If empty no filter will be applied. + lowpass_frequency: Union[list, int] + the list of the frequency for the low pass filter, in the same order as the index, if an integer is given + the same value will be used for each index + windows_length: Union[list, int] + The list of the window length for the moving average filter, + in the same order as the index, if an integer is given the same value will be used for each index + positions_from_inverse_kinematics: bool + If true use the result of precomputed inverse kinematics. + Note that the inverse kinematics must be computed before. + velocities_from_inverse_kinematics: bool + If true use the result of precomputed inverse kinematics. + Note that the inverse kinematics must be computed before. + Returns + ------- + np.ndarray + The generalized torque. + """ + tic = time.time() + if not biordb_package: + raise ModuleNotFoundError( + "Biorbd is not installed." + " Please install it via" + " 'conda install biorbd -cconda-forge' to use this function." + ) + if (positions_from_inverse_kinematics or velocities_from_inverse_kinematics) and len(self.kin_buffer) == 0: + raise ValueError("Inverse kinematics must be called before using kinematics results.") + states_init = [joint_positions, joint_velocities, joint_accelerations] + if isinstance(joint_positions, np.ndarray): + if len(joint_positions.shape) > 1 and joint_positions.shape[1] > 1: + raise RuntimeError( + "Data must be only for one frame," " please do a for loop if you need ID for more than one frame. " + ) + if positions_from_inverse_kinematics: + states_init[0] = self.kin_buffer[0][:, -1:] + if velocities_from_inverse_kinematics: + states_init[1] = self.kin_buffer[1][:, -1:] + if accelerations_from_inverse_kinematics: + states_init[2] = self.kin_buffer[2][:, -1:] + if not positions_from_inverse_kinematics and not joint_positions: + raise RuntimeError("Please provide at lease the joint position to compute the inverse dynamics.") + has_changed = self._state_idx_to_process != state_idx_to_process + self._state_idx_to_process = state_idx_to_process + if len(state_idx_to_process) != 0: + states_init = self._filter_states( + states_init, state_idx_to_process, windows_length, lowpass_frequency, has_changed + ) + + states = self._check_states(states_init) + tau = np.zeros((self.model.nbQ(), 1)) + # for i in range(tau.shape[1]): + if external_load is not None: + external_biorbd_loads = external_load.to_biorbd_loads(self.model) + tau[:, 0] = self.model.InverseDynamics( + states[0][:, -1], states[1][:, -1], states[2][:, -1], external_biorbd_loads + ).to_array() + else: + tau[:, 0] = self.model.InverseDynamics(states[0][:, -1], states[1][:, -1], states[2][:, -1]).to_array() + self.tau_buffer = ( + tau if len(self.tau_buffer) == 0 else np.append(self.tau_buffer[:, -self.data_windows + 1 :], tau, axis=1) + ) + self.process_time.append(time.time() - tic) + return self.tau_buffer.copy() + + def compute_static_optimization( + self, + q: np.ndarray = None, + q_dot: np.ndarray = None, + tau: np.ndarray = None, + ocp_solver: any = None, + compile_c_code: bool = True, + use_residual_torque: bool = True, + torque_tracking_as_objective: bool = True, + muscle_torque_dynamics_func: any = None, + scaling_factor: Union[list, tuple] = (1, 1), + muscle_track_idx: list = None, + emg: np.ndarray = None, + weight: dict = None, + x0: np.ndarray = None, + data_from_inverse_dynamics: bool = False, + solver_options: dict = None, + compile_only_first_call: bool = True, + print_optimization_status: bool = False, + ): + """ + Compute the static optimization using the model kinematics and a biorbd model type. + Parameters + ---------- + q: np.ndarray + The joint position for each model joints + q_dot: np.ndarray + The joint velocities for each model joints + tau: np.ndarray + The joint torques for each model joints + ocp_solver: AcadosOcpSolver + The acados solver to use, if none is provided a new one will be created + compile_c_code: bool + If true the c code will be compiled + use_residual_torque: bool + If true the residual torque will be used + torque_tracking_as_objective: bool + If true the torque tracking will be used as objective, otherwise it will be a constraint + muscle_torque_dynamics_func: ca.Function + The casadi function to use for the muscle torque dynamics (muscle activation -> muscle torque) + scaling_factor: list + The scaling factor to use for the muscle activation and muscle torque optimization variables + muscle_track_idx: list + The index of the muscles to track + emg: np.ndarray + The emg data to track for the muscle_track_idx provided + weight: dict + The weight to use for the objective function + x0: np.ndarray + The initial guess for the optimization + data_from_inverse_dynamics: bool + If true the data will be taken from the inverse dynamics (q, qdot, tau) instead of the provided data + solver_options: dict + The solver options to use for the acados solver + compile_only_first_call: bool + If true the c code will be compiled only for the first call + print_optimization_status: bool + If true the status of the optimization will be printed + Returns + ------- + tuple: + The optimal activation and torque for each muscles + """ + if muscle_track_idx and emg is None: + muscle_track_idx = None + if emg is not None and not muscle_track_idx: + raise RuntimeError("If you want to track muscles, you must provide the muscle index to track.") + + if not self.ca_model: + import biorbd_casadi as biorbd_ca + + self.ca_model = biorbd_ca.Model(self.model.path().absolutePath().to_string()) + if data_from_inverse_dynamics: + if len(self.id_state_buffer) == 0 or len(self.tau_buffer) == 0: + raise RuntimeError("You must have called the inverse dynamics before using the data from it.") + q = self.id_state_buffer[0][:, -1:] + q_dot = self.id_state_buffer[1][:, -1:] + tau = self.tau_buffer[:, -1:] + if q is None or q_dot is None or tau is None: + raise RuntimeError( + "Please provide q, q_dot and tau to compute the static optimization." + " Or use data from inverse dynamics." + ) + if ( + q.shape[0] != self.ca_model.nbQ() + or q_dot.shape[0] != self.ca_model.nbQ() + or tau.shape[0] != self.ca_model.nbQ() + ): + raise RuntimeError("The provided data must have the same number of dof as the model.") + if isinstance(q, np.ndarray): + if len(q.shape) > 1 and q.shape[1] > 1: + raise RuntimeError( + "Data must be only for one frame," " please do a for loop if you need ID for more than one frame. " + ) + if isinstance(q_dot, np.ndarray): + if len(q_dot.shape) > 1 and q_dot.shape[1] > 1: + raise RuntimeError( + "Data must be only for one frame," " please do a for loop if you need ID for more than one frame. " + ) + if isinstance(tau, np.ndarray): + if len(tau.shape) > 1 and tau.shape[1] > 1: + raise RuntimeError( + "Data must be only for one frame," " please do a for loop if you need ID for more than one frame. " + ) + + self.ocp_solver = ocp_solver if ocp_solver else self.ocp_solver + self.mjt_funct = muscle_torque_dynamics_func if muscle_torque_dynamics_func else self.mjt_funct + self.mjt_funct = self.mjt_funct if self.mjt_funct else _init_casadi_function(self.ca_model) + if not compile_c_code and not self.ocp_solver: + raise RuntimeError("You must provide a solver if you want avoid to compile c code.") + compile_c_code = not self.once_compile if compile_only_first_call else compile_c_code + if not self.ocp_solver or compile_c_code: + if not self.ocp: + self.ocp, self.weigh_list = _init_acados( + self.ca_model, + torque_tracking_as_objective, + self.mjt_funct, + use_residual_torque, + scaling_factor, + muscle_track_idx, + weight, + solver_options, + emg=emg, + ) + + self.ocp_solver = AcadosOcpSolver( + self.ocp, json_file=f"{self.ocp.model.name}.json", build=compile_c_code, generate=True + ) + self.once_compile = True + + target = np.zeros((len(self.weigh_list))) + # if emg is not None: + # target[np.where(np.array(self.weigh_list) == "tracking_emg")] = emg[:, 0] + self.ocp_solver = _update_solver( + self.ocp_solver, target, x0, q, q_dot, tau, emg=emg, torque_as_objective=torque_tracking_as_objective + ) + + self.ocp_solver.solve() + solution = self.ocp_solver.get(0, "x") + if print_optimization_status: + print("---------- QP Solver statistics ----------") + self.ocp_solver.print_statistics() + print("Residuals: ", self.ocp_solver.get_stats("residuals")) + print("Optimization status: ", self.ocp_solver.status, "\n") + muscle_activations = solution[: self.ca_model.nbMuscles() * q.shape[1]] / scaling_factor[0] + residual_torque = solution[self.ca_model.nbMuscles() * q.shape[1] :] / scaling_factor[1] + self.act_buffer = ( + np.append(self.act_buffer[:, -self.data_windows + 1 :], muscle_activations[:, np.newaxis], axis=1) + if len(self.act_buffer) != 0 + else muscle_activations[:, np.newaxis] + ) + self.res_tau_buffer = ( + np.append(self.res_tau_buffer[:, -self.data_windows + 1 :], residual_torque[:, np.newaxis], axis=1) + if len(self.res_tau_buffer) != 0 + else residual_torque[:, np.newaxis] + ) + return self.act_buffer.copy(), self.res_tau_buffer.copy() + + def compute_joint_reaction_load( + self, + q: np.ndarray = None, + qdot: np.ndarray = None, + qddot: np.ndarray = None, + muscle_activations: np.ndarray = None, + # residual_torques: np.ndarray = None, + act_from_static_optimisation: bool = False, + kinetics_from_inverse_dynamics: bool = False, + express_in_coordinate: str = None, + apply_on_segment: str = "all", + from_distal: bool = True, + application_point: Union[list, tuple] = None, + external_loads: any = None, + ): + if act_from_static_optimisation and len(self.act_buffer) == 0: + raise RuntimeError("You must compute muscle activation from static optimisation before using them.") + if (act_from_static_optimisation and muscle_activations is not None) or ( + not act_from_static_optimisation and muscle_activations is None + ): + raise RuntimeError( + "Please provide one muscle activation source. Either from static optimisation or " "from the user." + ) + if kinetics_from_inverse_dynamics and len(self.tau_buffer) == 0: + raise RuntimeError("You must compute joint kinetics from inverse dynamics before using them.") + if (kinetics_from_inverse_dynamics and np.sum((q, qdot, qddot)) is not None) or ( + not kinetics_from_inverse_dynamics and np.sum((q, qdot, qddot)) is None + ): + raise RuntimeError("Please provide one kinetics source. Either from inverse dynamics or " "from the user.") + add_idx = 0 if not from_distal else 1 + if self.model.segments()[-1].name().to_string() in apply_on_segment and from_distal: + raise RuntimeError( + "Can not give force from distal segment on the last segment." + "Please consider using directly the inverse dynamics." + ) + non_virtual_segments = [ + seg.name().to_string() for seg in self.model.segments() if seg.characteristics().mass() > 1e-7 + ] + + final_target_segments = ( + [ + non_virtual_segments[i + add_idx] + for i in range(len(non_virtual_segments) - 1) + if non_virtual_segments[i] in apply_on_segment + ] + if apply_on_segment != "all" + else non_virtual_segments + ) + if apply_on_segment != "all" and non_virtual_segments[-1] in apply_on_segment: + final_target_segments.append(non_virtual_segments[-1]) + express_in_coordinate = ( + [express_in_coordinate] if isinstance(express_in_coordinate, str) else express_in_coordinate + ) + application_point = [application_point] if not isinstance(application_point, list) else application_point + application_point = [[0, 0, 0]] * len(final_target_segments) if not application_point else application_point + if len(express_in_coordinate) != len(final_target_segments): + for coord in express_in_coordinate: + if coord not in final_target_segments: + raise RuntimeError( + "The segment provided is not a real segment but a virtual one. " + "Please provide a real segment to compute joint load." + ) + raise RuntimeError("You must provide the coordinate system to express your joint loads.") + if len(application_point) != len(final_target_segments): + raise RuntimeError("You must provide an application point for each wanted joint load.") + if isinstance(q, list): + q = np.array(q) + if isinstance(qdot, list): + qdot = np.array(qdot) + if isinstance(qddot, list): + qddot = np.array(qddot) + if len(q.shape) != 2: + q = q[:, np.newaxis] + if len(qdot.shape) != 2: + qdot = qdot[:, np.newaxis] + if len(qddot.shape) != 2: + qddot = qddot[:, np.newaxis] + if len(muscle_activations.shape) != 2: + muscle_activations = muscle_activations[:, np.newaxis] + if not self.model_all_dofs: + self.model_all_dofs, self.q_mapping, self.ordered_seg, self.ordered_idx = _create_new_model( + self.model, + final_target_segments, + ) + + q_tot = np.zeros((len(sum(self.ordered_idx, [])) + len(self.q_mapping), q.shape[1])) + q_dot_tot = np.zeros_like(q_tot) + q_ddot_tot = np.zeros_like(q_tot) + q_tot[self.q_mapping], q_dot_tot[self.q_mapping], q_ddot_tot[self.q_mapping] = q, qdot, qddot + q, qdot, qddot = q_tot, q_dot_tot, q_ddot_tot + all_trans = np.ndarray((len(final_target_segments), 3, q.shape[1])) + all_rot = np.ndarray((len(final_target_segments), 3, q.shape[1])) + muscle_activations = self.act_buffer if act_from_static_optimisation else muscle_activations + # res_torque = self.res_tau_buffer if act_from_static_optimisation else residual_torques + # b = bioviz.Viz(loaded_model=model_all_dofs) + # b.load_movement(q) + # b.exec() + for i in range(q.shape[1]): + tic = time.time() + all_global_jcs_old = [jcs.to_array() for jcs in self.model_all_dofs.allGlobalJCS(q[:, i])] + inv_all_global_jcs_new = [ + np.linalg.inv(jcs.to_array()) for jcs in self.model_all_dofs.allGlobalJCS(q[:, i]) + ] + translational_in_local, rotational_in_local = _compute_inverse_dynamics( + self.model_all_dofs, + q[:, i], + qdot[:, i], + qddot[:, i], + segment_names=self.ordered_seg, + segment_idx=self.ordered_idx, + external_loads=external_loads, + ) + if self.model_all_dofs.nbMuscles() != 0: + trans_muscle_actions, rot_muscle_actions = _compute_forces( + self.model_all_dofs, + q[:, i], + qdot[:, i], + muscle_activations[:, i], + segment_names=self.ordered_seg, + segment_idx=self.ordered_idx, + compound="muscle", + ) + + if self.model.nbLigaments() != 0: + raise RuntimeError("Ligaments are not yet implemented when computing joint loads.") + # rot_ligament_actions, trans_ligaments_actions = _compute_forces(model_all_dofs, q[:, i], qdot[:, i], + # muscle_activations[:, i], + # segment_names=ordered_seg, + # segment_idx=ordered_idx, + # compound="ligament") + if self.model.nbActuators() != 0: + raise RuntimeError("Actuators are not yet implemented when computing joint loads.") + if self.model.nbPassiveTorques() != 0: + raise RuntimeError("Passive torques are not yet implemented when computing joint loads.") + count = 0 + for k in range(len(self.ordered_seg)): + if self.ordered_seg[k] in final_target_segments: + all_trans[count, :, i] = translational_in_local[k] + all_rot[count, :, i] = rotational_in_local[k] + if self.model_all_dofs.nbMuscles() != 0: + all_trans[count, :, i] = np.sum((trans_muscle_actions[k], all_trans[count, :, i]), axis=0) + all_rot[count, :, i] = np.sum((rot_muscle_actions[k], all_rot[count, :, i]), axis=0) + if express_in_coordinate: + segment_idx = self.model_all_dofs.getBodyBiorbdId(final_target_segments[count]) + new_segment_idx = self.model_all_dofs.getBodyBiorbdId(express_in_coordinate[count]) + if new_segment_idx == -1: + raise RuntimeError( + f"The segment provided ({express_in_coordinate[count]}) does not exist." + " Please provide a real segment." + ) + all_trans[count, :, i], all_rot[count, :, i] = _express_in_new_coordinate( + all_trans[count, :, i], + all_rot[count, :, i], + application_point[count], + all_global_jcs_old[segment_idx], + inv_all_global_jcs_new[new_segment_idx], + ) + count += 1 + + # print("real_jrf_time:", time.time() - tic) + return np.concatenate((all_trans, all_rot), axis=0) + + def _check_states(self, states): + states_tmp = states.copy() + for s, state in enumerate(states_tmp): + states[s] = np.array(state) if isinstance(state, (tuple, list)) else state + states[s] = states[s][:, np.newaxis] if states[s] is not None and len(states[s].shape) != 2 else state + idx_to_compute_derivative = [i for i in range(len(states)) if states[i] is None] + all_shapes = [state.shape[1] for state in states if state is not None] + self.id_state_buffer = [None] * 3 if len(self.id_state_buffer) == 0 else self.id_state_buffer + for i in range(len(self.id_state_buffer)): + state_to_append = states[i] if states[i] is not None else np.zeros(states[0].shape) + if self.id_state_buffer[i] is None: + self.id_state_buffer[i] = state_to_append + else: + self.id_state_buffer[i] = np.append( + self.id_state_buffer[i][:, -self.data_windows + 1 :], state_to_append, axis=1 + ) + if all_shapes.count(all_shapes[0]) != len(all_shapes): + raise RuntimeError("Buffer and given data must have the same size.") + if len(idx_to_compute_derivative) >= 1: + self.id_state_buffer = self._compute_differential_state(idx_to_compute_derivative) + return self.id_state_buffer + + def _compute_differential_state(self, idx_to_compute_derivative): + if self.data_windows <= 2: + raise ValueError("Buffer size must be superior than 2.") + states = np.copy(self.id_state_buffer) + for i in range(1, len(states)): + if i in idx_to_compute_derivative: + derivative = ( + (states[i - 1][:, -2:-1] - states[i - 1][:, -1:]) / (2 / self.system_rate) + if states[i - 1].shape[1] > 1 + else np.zeros((states[i - 1].shape[0], 1)) + ) + self.id_state_buffer[i][:, -1:] = derivative + return self.id_state_buffer + + def _filter_states( + self, states, state_idx_to_process, windows_length=None, low_pass_frequency=None, has_changed=False + ): + self._rt_process_methods = ( + RealTimeProcessing(self.system_rate, self.data_windows) + if not self._rt_process_methods or has_changed + else self._rt_process_methods + ) + + states_to_process = None + for s, state in enumerate(states): + if s in state_idx_to_process: + if state is None: + raise RuntimeError("You must provide a values before filter it.") + states_to_process = ( + np.append(states_to_process, state, axis=0) if states_to_process is not None else state + ) + + states_proc = self._rt_process_methods.process_emg( + states_to_process, + moving_average=windows_length is not None, + low_pass_filter=low_pass_frequency is not None, + band_pass_filter=False, + centering=False, + absolute_value=False, + normalization=False, + moving_average_window=windows_length, + lpf_lcut=low_pass_frequency, + )[:, -1:] + for i, state in enumerate(states): + states[i] = ( + states_proc[i * state.shape[0] : (i + 1) * state.shape[0], :] if i in state_idx_to_process else state + ) + return states + def get_mean_process_time(self): """ Get the mean process time. @@ -171,3 +724,18 @@ def get_mean_process_time(self): The mean process time. """ return np.mean(self.process_time) + + def get_kinematics_from_ik(self): + return self.kin_buffer.copy() + + def get_tau_from_id(self): + return self.tau_buffer.copy() + + def get_filtered_kinematics_from_id(self): + return self.id_state_buffer.copy() + + def get_activation_from_so(self): + return self.act_buffer.copy() + + def get_jrf_from_external_load_analysis(self): + return self.jrf_buffer.copy() diff --git a/biosiglive/processing/msk_utils.py b/biosiglive/processing/msk_utils.py new file mode 100644 index 0000000..7f420fb --- /dev/null +++ b/biosiglive/processing/msk_utils.py @@ -0,0 +1,463 @@ +import numpy as np + +try: + import casadi as ca +except ModuleNotFoundError: + pass +try: + from acados_template import AcadosOcp, AcadosModel, AcadosOcpSolver + + acados_package = True +except ModuleNotFoundError: + acados_package = False +from scipy import linalg + +try: + import biorbd +except: + pass + + +def _set_solver_options(ocp, solver_options=None): + ocp.solver_options.qp_solver = "PARTIAL_CONDENSING_OSQP" # 'PARTIAL_CONDENSING_HPIPM' # FULL_CONDENSING_QPOASES + # PARTIAL_CONDENSING_HPIPM, FULL_CONDENSING_QPOASES, FULL_CONDENSING_HPIPM, + # PARTIAL_CONDENSING_QPDUNES, PARTIAL_CONDENSING_OSQP + ocp.solver_options.hessian_approx = "GAUSS_NEWTON" # GAUSS_NEWTON, EXACT + ocp.solver_options.integrator_type = "DISCRETE" + ocp.solver_options.sim_method_num_steps = 1 + ocp.solver_options.sim_method_num_stages = 1 + ocp.solver_options.sim_method_jac_reuse = 1 + ocp.solver_options.print_level = 0 + ocp.solver_options.tol = 1e-3 + ocp.solver_options.nlp_solver_type = "SQP_RTI" # SQP_RTI, SQP + # ocp.solver_options.levenberg_marquardt = 90.0 + ocp.solver_options.nlp_solver_max_iter = 1000 + ocp.solver_options.qp_solver_iter_max = 4000 + ocp.solver_options.qp_tol = 5e-5 + if solver_options: + for key, value in solver_options.items(): + setattr(ocp.solver_options, key, value) + return ocp + + +def _attribute_target_cost_and_constraints_function(model, names_J, tau, torque_tracking_as_objective): + constr = None + target = [] + for i in range(tau.shape[1]): + names_J.append(["act"] * model.nbMuscles()) + target.append([None] * model.nbMuscles()) + names_J.append(["pas_tau"] * model.nbQ()) + target.append([None] * model.nbQ()) + if torque_tracking_as_objective: + names_J.append(["tau"] * model.nbQ()) + for k in range(tau.shape[0]): + target.append([tau[k, i]]) + target = sum(target, []) + names_J = sum(names_J, []) + y_ref_start = [] + for idx in range(len(names_J)): + if target[idx] is not None: + y_ref_start.append(target[idx]) + else: + y_ref_start.append([0]) + return y_ref_start + + +def _init_acados( + model, + torque_tracking_as_objective, + mjt_func, + use_residual_torque, + scaling_factor, + muscle_track_idx, + weight, + solver_options, + emg=None, +): + q = ca.SX.sym("q", model.nbQ()) + qdot = ca.SX.sym("q", model.nbQ()) + tau = ca.SX.sym("tau", model.nbQ()) + emg_sym = None + if emg is not None: + emg_sym = ca.SX.sym("emg", emg.shape[0]) + x = ca.SX.sym("x", model.nbMuscles() * q.shape[1]) + if use_residual_torque: + pas_tau = ca.SX.sym("pas_tau", q.shape[0]) + else: + pas_tau = ca.SX.zeros(q.shape[0]) + lbx = (np.zeros((model.nbMuscles() * q.shape[1])) + 0.00001) * scaling_factor[0] + ubx = np.ones((model.nbMuscles() * q.shape[1])) * scaling_factor[0] + lh = tau[:, 0] + uh = tau[:, 0] + if use_residual_torque: + lb_torque = -np.ones((q.shape[0])) * 5 * scaling_factor[1] + ub_torque = np.ones((q.shape[0])) * 5 * scaling_factor[1] + lb_torque[:5] = -20 * scaling_factor[1] * np.ones((5)) + ub_torque[:5] = 20 * scaling_factor[1] * np.ones((5)) + # lb_torque[-1:] = -20 * scaling_factor[1] * np.ones((1)) + # ub_torque[-1:] = 20 * scaling_factor[1] * np.ones((1)) + # lb_torque = -20 * scaling_factor[1] * np.ones((q.shape[0])) + # ub_torque = 20 * scaling_factor[1] * np.ones((q.shape[0])) + lbx = np.hstack((lbx, lb_torque)) + ubx = np.hstack((ubx, ub_torque)) + ocp = AcadosOcp() + ocp.model = AcadosModel() + ocp = _set_solver_options(ocp, solver_options) + ocp, weigth_list = _init_cost_function( + model, + x, + mjt_func, + torque_tracking_as_objective, + pas_tau, + ocp, + q, + qdot, + tau, + emg_sym, + scaling_factor, + muscle_track_idx, + weight, + ) + x = ca.vertcat(x, pas_tau) + p = ca.vertcat(q, qdot) + if use_residual_torque: + if torque_tracking_as_objective: + p = ca.vertcat(p, tau) + if emg_sym is not None: + p = ca.vertcat(p, emg_sym) + ocp.model.p = p + ocp.model.disc_dyn_expr = x + ocp.dims.np = ocp.model.p.size()[0] + ocp.parameter_values = np.zeros((ocp.model.p.size()[0],)) + ocp.model.x = x + ocp.dims.nx = model.nbMuscles() * q.shape[1] + q.shape[0] + ocp.model.f_expl_expr = x + # ocp.model.f_impl_expr = x + ocp.model.xdot = x + # ocp.model.con_h_expr = constr + + ocp.dims.nbx = ocp.dims.nx + ocp.model.u = ca.SX.sym("u", 0) + ocp.model.name = "test_acados_SO" + T = 1 + N = 1 + ocp.dims.N = N + ocp.solver_options.tf = T + ocp.constraints.idxbx_0 = np.array(range(ocp.dims.nx)) + if not torque_tracking_as_objective: + ocp.dims.nh = tau.shape[0] + ocp.constraints.lh = np.zeros((tau.shape[0],)) + ocp.constraints.uh = np.zeros((tau.shape[0],)) + ocp.constraints.lbx_0 = lbx + ocp.constraints.ubx_0 = ubx + return ocp, weigth_list + + +def _init_cost_function( + model, + x, + ca_function, + torque_tracking_as_objective, + pas_tau, + ocp, + q, + qdot, + tau, + emg, + scaling_factor, + muscle_track_idx, + weights, +): + if not weights: + weights = {"tau": 1, "act": 1, "tracking_emg": 1, "pas_tau": 1} + # emg = np.ones((len(muscle_track_idx), q.shape[1])) * 0.7 + names_J = [] + J = None + constr = None + for i in range(q.shape[1]): + mus_tau = ca_function( + x[i * model.nbMuscles() : (i + 1) * model.nbMuscles()] / scaling_factor[0], q[:, i], qdot[:, i] + ) + + for m in range(model.nbMuscles()): + if emg is not None and (muscle_track_idx and m in muscle_track_idx): + idx = muscle_track_idx.index(m) + if J is None: + J = ( + (x[i * model.nbMuscles() + m : i * model.nbMuscles() + m + 1]) + - ca.SX(emg[idx, i]) * scaling_factor[0] + ) ** 2 + names_J.append(["tracking_emg"]) + # J = ca.vertcat(J, x[i * model.nbMuscles() + m: i * model.nbMuscles() + m + 1] ** 2) + # names_J.append(["act"]) + else: + J = ca.vertcat( + J, + ( + (x[i * model.nbMuscles() + m : i * model.nbMuscles() + m + 1]) + - ca.SX(emg[idx, i]) * scaling_factor[0] + ) + ** 2, + ) + names_J.append(["tracking_emg"]) + # J = ca.vertcat(J, x[i * model.nbMuscles() + m: i * model.nbMuscles() + m + 1] ** 2) + # names_J.append(["act"]) + else: + if J is None: + J = x[i * model.nbMuscles() + m : i * model.nbMuscles() + m + 1] ** 2 + names_J.append(["act"]) + else: + J = ca.vertcat(J, x[i * model.nbMuscles() + m : i * model.nbMuscles() + m + 1] ** 2) + names_J.append(["act"]) + + J = ca.vertcat(J, pas_tau**2) + names_J.append(["pas_tau"] * model.nbQ()) + if torque_tracking_as_objective: + J = ca.vertcat(J, (tau - (mus_tau + pas_tau / scaling_factor[1])) ** 2) + names_J.append(["tau"] * model.nbQ()) + else: + if constr is None: + constr = mus_tau + pas_tau + else: + constr = ca.vertcat(constr, (mus_tau + pas_tau)) + + names_J = sum(names_J, []) + ocp.cost.cost_type_0 = "NONLINEAR_LS" + w = [] + for idx in range(J.numel()): + w.append(weights[names_J[idx]]) + ocp.cost.W_0 = linalg.block_diag(np.diag(w)) + ocp.cost.yref_0 = np.zeros((ocp.cost.W_0.shape[0],)) + ocp.dims.ny_0 = ocp.cost.W_0.shape[0] + ocp.model.cost_y_expr_0 = J.reshape((-1, 1)) + if not torque_tracking_as_objective: + ocp.constraints.constr_type = "BGH" + ocp.model.con_h_expr = constr + return ocp, names_J + + +def _init_casadi_function(model): + q_sym = ca.MX.sym("q", model.nbQ()) + qdot_sym = ca.MX.sym("qdot", model.nbQ()) + muscle_activations = ca.MX.sym("muscle_activation", model.nbMuscles()) + + def muscle_joint_torque(activations_fct, q_fct, qdot_fct) -> ca.MX: + muscles_states = model.stateSet() + for k in range(model.nbMuscles()): + muscles_states[k].setActivation(activations_fct[k]) + return model.muscularJointTorque(muscles_states, q_fct, qdot_fct).to_mx() + + mjt = muscle_joint_torque(muscle_activations, q_sym, qdot_sym) + mjt_func = ca.Function("mjt_func", [muscle_activations, q_sym, qdot_sym], [mjt]).expand() + return mjt_func + + +def _update_solver(ocp_solver, target, x0, q, qdot, tau=None, torque_as_objective=True, emg=None): + # update initial guess + # if x0 is not None: + # ocp_solver.set(0, "x", x0) + # update targets + # ocp_solver.cost_set(0, "yref", target) + if not torque_as_objective: + ocp_solver.constraints_set(0, "lh", tau[:, 0]) + ocp_solver.constraints_set(0, "uh", tau[:, 0]) + if emg is not None: + ocp_solver.set(0, "p", np.vstack((q, qdot, emg))) + else: + ocp_solver.set(0, "p", np.vstack((q, qdot))) + + else: + if emg is not None: + ocp_solver.set(0, "p", np.vstack((q, qdot, tau, emg))) + else: + ocp_solver.set(0, "p", np.vstack((q, qdot, tau))) + return ocp_solver + + +def _create_new_model(model, segments_names): + model_empty = biorbd.Model() + ordered_dof_prox_to_dist = [] + ordered_dof_prox_to_dist_idx = [] + q_non_zero_idx = [] + nb_dof = 0 + for i in range(model.nbSegment()): + last_parent = None + # if i == 0: + # ordered_dof_prox_to_dist = [model.segment(i).name().to_string()] + if model.segment(i).name().to_string() in segments_names: + rot = model.segment(i).seqR().to_string() + trans = model.segment(i).seqT().to_string() + if len(rot + trans) != 0: + for j in range(len(rot + trans)): + q_non_zero_idx.append(nb_dof + j) + model_empty = __add_segment_to_model( + model_empty, model.segment(i), rot, trans, name=model.segment(i).name().to_string() + "_init_dof" + ) + last_parent = model.segment(i).name().to_string() + "_init_dof" + nb_dof += len(rot + trans) + rot = trans = "xyz" + ordered_dof_prox_to_dist.append(model.segment(i).name().to_string()) + ordered_dof_prox_to_dist_idx.append([nb_dof + k for k in range(6)]) + nb_dof += 6 + model_empty = __add_segment_to_model( + model_empty, model.segment(i), rot, trans, parent_str=last_parent, RT=biorbd.RotoTrans() + ) + else: + rot = model.segment(i).seqR().to_string() + trans = model.segment(i).seqT().to_string() + for j in range(len(rot + trans)): + q_non_zero_idx.append(nb_dof + j) + nb_dof += len(rot + trans) + model_empty = __add_segment_to_model(model_empty, model.segment(i), rot, trans) + for i, group in enumerate(model.muscleGroups()): + name, insertion, origin = group.name().to_string(), group.insertion().to_string(), group.origin().to_string() + model_empty.addMuscleGroup(name, origin, insertion) + for m in range(group.nbMuscles()): + model_empty.muscleGroups()[-1].addMuscle(model.muscleGroup(i).muscle(m)) + # for i, group in enumerate(model.ligaments()): + # name, insertion, origin = group.name().to_string(), group.insertion().to_string(), group.origin().to_string() + # model_empty.addMuscleGroup(name, origin, insertion) + # for m in range(group.nbMuscles()): + # model_empty.muscleGroups()[-1].addMuscle(model.muscleGroup(i).muscle(m)) + # if write_bioMod: + # biorbd.Writer.writeModel(model_empty, "model_simplified.bioMod") + return model_empty, q_non_zero_idx, ordered_dof_prox_to_dist, ordered_dof_prox_to_dist_idx + + +def _compute_inverse_dynamics(model, q, qdot, qddot, segment_names, segment_idx, external_loads=None): + external_biorbd_loads = None + if external_loads: + external_biorbd_loads = external_loads.to_biorbd_loads(model) + Tau = model.InverseDynamics(q, qdot, qddot, external_biorbd_loads).to_array() + translational_in_local = [] + rotationnal_in_local = [] + for j in range(len(segment_names)): + translational_in_local.append(-Tau[segment_idx[j][:3]]) + rotationnal_in_local.append(-Tau[segment_idx[j][3:]]) + return translational_in_local, rotationnal_in_local + + +def _express_in_new_coordinate(trans, rot, new_application, all_global_jcs_old, inv_all_global_jcs_new): + def get_homogeneous_vector(vector): + if isinstance(vector, list): + vector = np.array(vector) + return np.append(vector, np.ones(1)) + + global_trans = np.dot(all_global_jcs_old, get_homogeneous_vector(trans)) + global_rot = np.dot(all_global_jcs_old, get_homogeneous_vector(rot)) + application_point = [0, 0, 0, 1] + application_point_global = all_global_jcs_old @ application_point + trans_new_coordinate = np.dot(inv_all_global_jcs_new, global_trans)[:3] + rot_new_coordinate = np.dot(inv_all_global_jcs_new, global_rot)[:3] + application_point_local = (inv_all_global_jcs_new @ application_point_global)[:3] + final_rotation = rot_new_coordinate + np.cross((application_point_local - new_application), trans_new_coordinate) + return trans_new_coordinate, final_rotation + + +def __add_segment_to_model(model, segment, rot, trans, name=None, RT=None, parent_str=None): + (name_tmp, parent_str_tmp, rot, trans, QRanges, QDotRanges, QDDotRanges, characteristics, RT_tmp) = ( + segment.name().to_string(), + segment.parent().to_string(), + rot, + trans, + [biorbd.Range(-3, 3)] * (len(rot) + len(trans)), + [biorbd.Range(-3 * 10, 3 * 10)] * (len(rot) + len(trans)), + [biorbd.Range(-3 * 100, 3 * 100)] * (len(rot) + len(trans)), + segment.characteristics(), + segment.localJCS(), + ) + name = name if name else name_tmp + RT = RT if RT else RT_tmp + parent_str = parent_str if parent_str else parent_str_tmp + model.AddSegment(name, parent_str, trans, rot, QRanges, QDotRanges, QDDotRanges, characteristics, RT) + return model + + +def _compute_forces(model, q, qdot, act, segment_names, segment_idx, compound="muscle"): + moment_arm = None + if compound == "muscle": + muscles_states = model.stateSet() + for m in range(model.nbMuscles()): + muscles_states[m].setActivation(act[m]) + moment_arm = model.muscularJointTorque(muscles_states, q, qdot).to_array() + if compound == "ligament": + moment_arm = model.ligamentJointTorque(q, qdot).to_array() + idx_segment = [] + for k in range(model.nbSegment()): + if model.segment(k).name().to_string() in segment_names: + idx_segment.append(k) + translational_in_local = [] + rotationnal_in_local = [] + for i in range(len(idx_segment)): + translational_in_local.append(moment_arm[segment_idx[i][:3]]) + # translational_in_global.append(np.dot(model.globalJCS(idx_segment[i]).to_array(), + # np.array(local_forces.tolist() + [1]))) + rotationnal_in_local.append(moment_arm[segment_idx[i][3:]]) + return translational_in_local, rotationnal_in_local + + +class ExternalLoad: + def __init__(self, point_of_application, applied_on_body, force, express_in_coordinate: str = "ground", name=""): + self.point_of_application = point_of_application + self.applied_on_body = applied_on_body + self.express_in_coordinate = express_in_coordinate + self.name = name + self.force = force + if isinstance(point_of_application, list): + self.point_of_application = np.array(point_of_application) + if self.point_of_application.shape[0] != 3: + raise ValueError("The point of application must be a 3xn vector") + if len(self.point_of_application.shape) != 1: + raise ValueError("The point of application must be a one dimensional array") + if isinstance(force, list): + self.force = np.array(force) + if self.force.shape[0] != 6: + raise ValueError("The force must be a 6xn vector") + if applied_on_body != express_in_coordinate and express_in_coordinate != "ground": + raise NotImplementedError( + "You can only either express the force in the ground or" + " on the body where the force is applied." + f"You have {applied_on_body} and {express_in_coordinate}" + ) + + +class ExternalLoads: + def __init__(self): + self.external_loads = [] + + def add_external_load(self, point_of_application, applied_on_body, express_in_coordinate, load, name=None): + self.external_loads.append( + ExternalLoad(point_of_application, applied_on_body, load, express_in_coordinate, name) + ) + + def update_external_load_value(self, value: np.ndarray, name: str = None, idx: int = None): + external_load = self.get_external_load(idx, name) + if isinstance(value, list): + value = np.array(value) + if value.shape[0] != 6: + raise ValueError(f"The load must be a 6xn vector." f" You provided a vector of shape {value.shape}.") + external_load.force = value + + def get_external_load(self, idx: int = None, name: str = None): + if (not idx and not name) or (idx and name): + raise RuntimeError("Please provide either the force index or name.") + if idx: + return self.external_loads[idx] + if name: + return [ext_load for ext_load in self.external_loads if ext_load.name == name][0] + + def to_biorbd_loads(self, model): + idx = list(range(len(self.external_loads))) + return self.to_biorbd_load(idx=idx, model=model) + + def to_biorbd_load(self, model, name: (str, list) = None, idx: (int, list) = None): + if (not idx and not name) or (idx and name): + raise RuntimeError("Please provide either the force index or name.") + ext_load = model.externalForceSet() + for i, load in enumerate(self.external_loads): + if i in idx or load.name == name: + if load.express_in_coordinate == "ground": + ext_load.add(load.applied_on_body, load.force[:, 0], load.point_of_application) + else: + ext_load.addInSegmentReferenceFrame(load.applied_on_body, load.force, load.point_of_application) + return ext_load diff --git a/biosiglive/streaming/async_client.py b/biosiglive/streaming/async_client.py new file mode 100644 index 0000000..457bb5b --- /dev/null +++ b/biosiglive/streaming/async_client.py @@ -0,0 +1,94 @@ +import asyncio +import struct +import numpy as np + + +class AsyncTCPClient: + """ + Async TCP client for sending NumPy arrays. + + Parameters + ------------ + host : str + port : int + + Returns + -------- + None + """ + + def __init__(self, host: str = "127.0.0.1", port: int = 5000, timeout: float = 5.0): + self.host = host + self.port = port + self.timeout = timeout + self.writer = None + + async def connect(self) -> None: + """ + Connect to server. + + Parameters + ------------ + None + + Returns + -------- + None + """ + try: + _, self.writer = await asyncio.wait_for(asyncio.open_connection(self.host, self.port), timeout=self.timeout) + except asyncio.TimeoutError: + print(f"Connection to {self.host}:{self.port} timed out after {self.timeout} seconds.") + + async def send_array(self, array: np.ndarray, sample_time: np.ndarray = None) -> None: + """ + Send NumPy array. + + Parameters + ------------ + array : np.ndarray + sample_time : np.ndarray, optional + The timestamps corresponding to each sample in the array. If provided, the timestamps will be sent along with the array. The server should be able to handle this case. + Returns + -------- + None + """ + if self.writer is None: + raise RuntimeError("Client not connected") + if array.dtype != np.float64: + array = array.astype(np.float64) + + if sample_time is not None: + if sample_time.dtype != np.float64: + sample_time = sample_time.astype(np.float64) + if len(sample_time) != array.shape[-1]: + return + array = np.vstack((sample_time, array)) + + payload = array.tobytes() + + shape = array.shape + ndim = len(shape) + + is_time = struct.pack("!?", sample_time is not None) + header = struct.pack("!II", len(payload), ndim) + shape_bytes = struct.pack(f"!{ndim}I", *shape) + + self.writer.write(is_time + header + shape_bytes + payload) + await self.writer.drain() + + async def close(self) -> None: + """ + Close connection. + + Parameters + ------------ + None + + Returns + -------- + None + """ + if self.writer: + self.writer.close() + await self.writer.wait_closed() diff --git a/biosiglive/streaming/async_server.py b/biosiglive/streaming/async_server.py new file mode 100644 index 0000000..6ca6d77 --- /dev/null +++ b/biosiglive/streaming/async_server.py @@ -0,0 +1,164 @@ +import asyncio +import struct +import numpy as np +import time +from .utils import CircularBuffer, RollingBuffer + + +class AsyncTCPServer: + """ + Async TCP server for receiving NumPy arrays. + + Parameters + ------------ + host : str + port : int + + Returns + -------- + None + """ + + def __init__(self, host: str = "127.0.0.1", port: int = 5000, buffer_length: int = 1000): + self.host = host + self.port = port + self.server = None + self.buffer = None + self.buffer_length = buffer_length + self.dt = 0.01 + + def init_buffer(self, n_channels, dt=None): + if dt is not None: + self.dt = dt + self.buffer = CircularBuffer(n_channels, self.buffer_length, dtype=np.float64, dt=self.dt) + # self.buffer = RollingBuffer(n_channels, self.buffer_length) + + async def start(self, task: callable = None) -> None: + """ + Start the TCP server. + + Parameters + ------------ + task : callable + Optional task to run when the server receive new data. The task should be a function that takes a NumPy array as input. Becarefull, the task will be called for each new data received, so it should be fast to avoid blocking the server. + + Returns + -------- + None + """ + self.task = task + self.server = await asyncio.start_server(self.handle_client, self.host, self.port) + # print(f"Server started on {self.host}:{self.port}") + + async with self.server: + await self.server.serve_forever() + + async def handle_client( + self, + reader: asyncio.StreamReader, + writer: asyncio.StreamWriter, + ) -> None: + """ + Handle incoming client connection. + + Parameters + ------------ + reader : asyncio.StreamReader + writer : asyncio.StreamWriter + + Returns + -------- + None + """ + addr = writer.get_extra_info("peername") + try: + while True: + data, t = await self._read_array(reader) + if data is None: + break + if self.buffer is None: + self.init_buffer(data.shape[0]) + + # if data.shape[0] != self.buffer.shape[0]: + # raise ValueError(f"Received data with shape {data.shape}, but expected {self.buffer.shape}.") + # if np.any(np.diff(t) <= 0): + # print(" intra-chunk non-monotonic") + self.buffer.append(data, t) + self.task(data, t) if self.task else None + + except asyncio.IncompleteReadError: + pass + + finally: + # print(f"Client disconnected: {addr}") + writer.close() + # await writer.wait_closed() + + async def _read_array(self, reader: asyncio.StreamReader) -> tuple: + """ + Read a NumPy array from stream. + + Parameters + ------------ + reader : asyncio.StreamReader + + Returns + -------- + tuple + """ + is_time = await reader.readexactly(1) + is_time = struct.unpack("!?", is_time)[0] + + header = await reader.readexactly(8) + size, ndim = struct.unpack("!II", header) + + shape_bytes = await reader.readexactly(4 * ndim) + shape = struct.unpack(f"!{ndim}I", shape_bytes) + + payload = await reader.readexactly(size) + array = np.frombuffer(payload, dtype=np.float64).reshape(shape) + + if is_time: + t = array[0, :] + array = array[1:, :] + if self.dt is None: + self.dt = t[1] - t[0] + self.tol = self.dt * 1.5 + else: + t = None + + return array, t + + async def stop(self) -> None: + """ + Stop the TCP server. + + Parameters + ------------ + None + + Returns + -------- + None + """ + if self.server: + self.server.close() + await self.server.wait_closed() + # print("Server stopped") + + async def get_data(self) -> tuple: + """ + Get the current data and timestamps in the buffer. + + Parameters + ------------ + None + + Returns + -------- + tuple + """ + if self.buffer is not None: + return self.buffer.get() + else: + raise ValueError("Buffer is empty. No data received yet.") diff --git a/biosiglive/streaming/server.py b/biosiglive/streaming/server.py index 2a9cfc7..7c75ea4 100644 --- a/biosiglive/streaming/server.py +++ b/biosiglive/streaming/server.py @@ -1,6 +1,7 @@ """ This file contains a wrapper for the socket server to send data to the server. """ + import socket import numpy as np import struct diff --git a/biosiglive/streaming/utils.py b/biosiglive/streaming/utils.py index c894701..4fa19ca 100644 --- a/biosiglive/streaming/utils.py +++ b/biosiglive/streaming/utils.py @@ -1,15 +1,160 @@ import numpy as np -def dic_merger(dic_to_merge, new_dic=None): +class CircularBuffer: + def __init__(self, n, W, dtype=np.float64, time_dtype=np.float64, dt=None): + self.W = W + self.shape = (n, W) + self.ring = np.zeros(self.shape, dtype=dtype) + + t = np.arange(W) * dt if dt is not None else np.arange(W) + self.ring_t = None + self.has_last = False + self.total_samples = 0 + self.idx = 0 + self.full = False + self.version = 0 + + if dt is not None: + self.dt = dt + self.tol = dt * 1.5 + + @property + def empty(self): + return self.idx == 0 and not self.full + + def append( + self, + x: np.ndarray, + t: np.ndarray = None, + fill_discontinuous: bool = False, + fill_mode="single", + row_idx=None, + dt=None, + ): + """ + Parameters + ---------- + x : np.ndarray + The data to append. Shape: (n, w) + t : np.ndarray, optional + The timestamps corresponding to the data. Shape: (w,) + fill_discontinuous : bool, optional + Whether to fill discontinuities in the data with one NaN value. Mostly for use with pyqtgraph to plot gaps. Default is False. + fill_mode : str, optional + The mode for filling discontinuities. Default is "single". Possible values are "single" (fill with a single NaN value) and "full" (fill all missing data). Only used if fill_discontinuous is True. + + """ + row_idx = slice(None, None, None) if row_idx is None else row_idx + + self._append(x, t, row_idx) + + def _append(self, data, t=None, row_idx=None): + """ + x shape: (n, w) + """ + n = data.shape[1] + start_version = self.version + + k = self.idx + end = (k + n) % self.W + + if end <= k: + split = self.W - k + self.ring[:, k:] = data[:, :split] + self.ring[:, :end] = data[:, split:] + if t is not None: + if self.ring_t is None: + self.ring_t = np.zeros(self.W, dtype=np.float64) + self.ring_t[k:] = t[:split] + self.ring_t[:end] = t[split:] + + else: + self.ring[:, k:end] = data + if t is not None: + if self.ring_t is None: + self.ring_t = np.zeros(self.W, dtype=np.float64) + self.ring_t[k:end] = t + + self.idx = end % self.W + self.total_samples += n + self.full |= self.total_samples >= self.W + self.version = start_version + 1 + + def get(self): + while True: + v1 = self.version + + # read all mutable state immediately after v1 + k = self.idx + size = min(self.total_samples, self.W) + + start = (k - size) % self.W + + if start < k: + data = self.ring[:, start:k].copy() + t = self.ring_t[start:k].copy() if self.ring_t is not None else None + else: + data = np.concatenate((self.ring[:, start:], self.ring[:, :k]), axis=-1).copy() + t = np.concatenate((self.ring_t[start:], self.ring_t[:k])).copy() if self.ring_t is not None else None + + v2 = self.version + if v1 == v2: + return data, t # ← was silently dropping t + + def append_row(self, x, t=None, fill_discontinuous=False, fill_mode="single", row_idx=None): + """Append a row to the buffer. The row_idx parameter specifies which row to append to. The rest of the rows will be filled with NaN values. + Parameters + ---------- + x : np.ndarray + The data to append. Shape: (w,) + t : np.ndarray, optional + The timestamps corresponding to the data. Shape: (w,) + fill_discontinuous : bool, optional + Whether to fill discontinuities in the data with one NaN value. Mostly for use with pyqtgraph to plot gaps. Default is False. + fill_mode : str, optional + The mode for filling discontinuities. Default is "single". Possible values are "single" (fill with a single NaN value) and "full" (fill all missing data). Only used if fill_discontinuous is True. + row_idx : int, optional + The index of the row to append to. Default is 0. + """ + row_idx = slice(None, None, None) if row_idx is None else row_idx + + self.append(x[np.newaxis, :], t, fill_discontinuous, fill_mode) + + +class RollingBuffer: + def __init__(self, n, W): + self.data = np.zeros((n, W), dtype=np.float32) + self.time = np.zeros(W, dtype=np.float64) + + def append(self, x, t): + w = x.shape[-1] + + self.data = np.roll(self.data, -w, axis=1) + self.data[:, -w:] = x + + self.time = np.roll(self.time, -w) + self.time[-w:] = t + + def get_time(self, len=None): + return self.time + + def get_data(self, len=None): + return self.data + + def get(self, len=None): + return self.get_data(len), self.get_time(len) + + +def dic_merger(dic_to_merge: dict, new_dic: dict = None) -> dict: """Merge two dictionaries. Parameters ---------- dic_to_merge : dict - Dictionary to merge. + Existing dictionary to merge. new_dic : dict - Dictionary to merge with. + Temporary dictionary to merge with. Returns ------- @@ -21,14 +166,32 @@ def dic_merger(dic_to_merge, new_dic=None): new_dic = dic_to_merge else: for key in dic_to_merge.keys(): + if dic_to_merge[key] is None: + dic_to_merge[key] = [dic_to_merge[key]] + if new_dic[key] is None: + new_dic[key] = [new_dic[key]] + if isinstance(new_dic[key], (int, float, str)): + new_dic[key] = [new_dic[key]] + if isinstance(dic_to_merge[key], (int, float, str)): + dic_to_merge[key] = [dic_to_merge[key]] if isinstance(dic_to_merge[key], dict): - new_dic[key] = dic_merger(dic_to_merge[key], new_dic[key]) + if len(new_dic[key].keys()) == 0: + new_dic[key] = dic_to_merge[key] + else: + new_dic[key] = dic_merger(dic_to_merge[key], new_dic[key]) elif isinstance(dic_to_merge[key], list): new_dic[key] = dic_to_merge[key] + new_dic[key] elif isinstance(dic_to_merge[key], np.ndarray): - new_dic[key] = np.append(dic_to_merge[key], new_dic[key], axis=0) - elif isinstance(dic_to_merge[key], int): - if isinstance(new_dic[key], int): - new_dic[key] = [new_dic[key]] - new_dic[key] = [new_dic[key]] + [dic_to_merge[key]] + if not isinstance(new_dic[key], np.ndarray): + new_dic[key] = np.array(new_dic[key]) + if len(new_dic[key].shape) == 1: + new_dic[key] = new_dic[key][:, np.newaxis] + if len(dic_to_merge[key].shape) == 1: + dic_to_merge[key] = dic_to_merge[key][:, np.newaxis] + new_dic[key] = np.append(dic_to_merge[key], new_dic[key], axis=-1) + else: + raise ValueError("Type not supported") + for key in new_dic.keys(): + if key not in dic_to_merge.keys(): + new_dic[key] = new_dic[key] return new_dic diff --git a/environment.yml b/environment.yml index 738b5d8..d3aa954 100644 --- a/environment.yml +++ b/environment.yml @@ -10,4 +10,5 @@ dependencies: - scipy - biorbd - bioviz -- pyqtgraph \ No newline at end of file +- pyqtgraph +- PyQt6 \ No newline at end of file diff --git a/examples/EMG_streaming.py b/examples/EMG_streaming.py index 731b69d..a8d2791 100644 --- a/examples/EMG_streaming.py +++ b/examples/EMG_streaming.py @@ -19,10 +19,13 @@ In this function you can pass a method if you want to use a different method than the default one and every needed argument for that function as well. """ +from biosiglive.enums import DeviceType from custom_interface import MyInterface from biosiglive.gui.plot import LivePlot from biosiglive import ( ViconClient, + PytrignoClient, + save, RealTimeProcessingMethod, PlotType, ) @@ -30,24 +33,24 @@ if __name__ == "__main__": - try_offline = True + try_offline = False output_file_path = "trial_x.bio" if try_offline: interface = MyInterface(system_rate=100, data_path="abd.bio") else: - interface = ViconClient(ip="localhost", system_rate=100) + interface = PytrignoClient(ip="127.0.0.1") # or interface = ViconClient(ip="localhost", system_rate=100) - n_electrodes = 4 + n_electrodes = 2 muscle_names = [ "Pectoralis major", "Deltoid anterior", - "Deltoid medial", - "Deltoid posterior", + # "Deltoid medial", + # "Deltoid posterior", ] interface.add_device( nb_channels=n_electrodes, - device_type="emg", + device_type=DeviceType.Emg, name="emg", rate=2000, device_data_file_key="emg", @@ -67,6 +70,7 @@ time_to_sleep = 1 / 100 count = 0 + file_path = "data_test.bio" while True: tic = time() raw_emg = interface.get_device_data(device_name="emg") @@ -74,7 +78,9 @@ emg_plot.update(emg_proc[:, -1:]) emg_raw_plot.update(raw_emg) count += 1 - loop_time = time() - tic - real_time_to_sleep = time_to_sleep - loop_time - if real_time_to_sleep > 0: - sleep(real_time_to_sleep) + # save({'emg_raw': raw_emg, 'proc_emg': emg_proc}, file_path, add_data=True) + # if try_offline: + # loop_time = time() - tic + # real_time_to_sleep = time_to_sleep - loop_time + # if real_time_to_sleep > 0: + # sleep(real_time_to_sleep) diff --git a/examples/async_client.py b/examples/async_client.py new file mode 100644 index 0000000..dd19a36 --- /dev/null +++ b/examples/async_client.py @@ -0,0 +1,22 @@ +from biosiglive import AsyncTCPClient + +import asyncio +import numpy as np + + +async def main(): + client = AsyncTCPClient("127.0.0.1", 12345) + await client.connect() + dt = 0.01 + t0 = 0 + + while True: + t = t0 + np.arange(30) * dt + t0 += 30 * dt + data = np.random.randn(10, 30).astype(np.float64) + await client.send_array(data, sample_time=t) + await asyncio.sleep(0.01) # ~100 Hz + + +if __name__ == "__main__": + asyncio.run(main()) diff --git a/examples/async_server.py b/examples/async_server.py new file mode 100644 index 0000000..1af5d1a --- /dev/null +++ b/examples/async_server.py @@ -0,0 +1,33 @@ +import numpy as np + +from biosiglive import AsyncTCPServer +from biosiglive.streaming.utils import CircularBuffer +import asyncio +import threading +import time + +if __name__ == "__main__": + server = AsyncTCPServer("127.0.0.1", 12345) + + def task(data, t): + pass + + threading.Thread(target=asyncio.run, args=(server.start(task=task),)).start() + data_buffer = CircularBuffer(10, 1000) + + last_idx = None + while True: + time.sleep(0.05) + if server.buffer is None: + continue + data, t = server.buffer.get_view() + current_buff_idx = server.buffer.total_samples + if last_idx: + n_new_data = current_buff_idx - last_idx + if n_new_data == 0: + continue + data = data[:, -n_new_data:] + t = t[-n_new_data:] + last_idx = current_buff_idx + data_buffer.append(data, t) + data_n, t_n = data_buffer.get_view() diff --git a/examples/cadence_from_treadmill.py b/examples/cadence_from_treadmill.py index ef34bfa..c800bec 100644 --- a/examples/cadence_from_treadmill.py +++ b/examples/cadence_from_treadmill.py @@ -4,6 +4,7 @@ The data can be plotted in real time at each loop, please see the live_plot.py example. """ + from biosiglive.interfaces.vicon_interface import ViconClient from time import time, sleep from custom_interface import MyInterface diff --git a/examples/custom_interface.py b/examples/custom_interface.py index efd504e..b611032 100644 --- a/examples/custom_interface.py +++ b/examples/custom_interface.py @@ -38,14 +38,11 @@ def add_device( name: str = None, data_buffer_size: int = None, rate: float = 2000, - device_range: tuple = None, device_data_file_key: str = None, processing_method: Union[RealTimeProcessingMethod, OfflineProcessingMethod] = None, **process_kwargs, ): - device_tmp = self._add_device( - nb_channels, device_type, name, rate, device_range, processing_method, **process_kwargs - ) + device_tmp = self._add_device(nb_channels, device_type, name, rate, processing_method, **process_kwargs) if data_buffer_size: device_tmp.data_window = data_buffer_size self.devices.append(device_tmp) diff --git a/examples/inverse_kinematics_from_mocap.py b/examples/inverse_kinematics_from_mocap.py index 469cf49..5125d2d 100644 --- a/examples/inverse_kinematics_from_mocap.py +++ b/examples/inverse_kinematics_from_mocap.py @@ -6,6 +6,7 @@ The processing method will return the angle and velocity of the model joint. You can display the result using the skeleton plot which uses the bioviz library. In the skeleton plot initialization, you can specify the arguments belonging to the bioviz.Viz function. For more information about the arguments, please refer to the bioviz documentation. Be aware that the skeleton plot may take some time and slow down the loop. If you want to display the data in real time, consider using a lightweight *.vtp file inside the model. """ + import time import numpy as np diff --git a/examples/live_mvc.py b/examples/live_mvc.py index eab9698..305bbfc 100644 --- a/examples/live_mvc.py +++ b/examples/live_mvc.py @@ -7,6 +7,7 @@ You can plot the data (raw, processed or both) at the end of the processing. Then you can continue with other trials and do the same commands until you decide to stop the program at the end of a trial. In case of malfunction, a temporary file is saved after each trial with all previously recorded data (raw and processed). At the end of all trials, the MVC will be calculated by sorting the data and taking the highest values (for the specified MVC_window). You can then save the data to a file or simply terminate the program which will return a list of MVCs of the length of the number of muscles. """ + from biosiglive import ComputeMvc, InterfaceType from custom_interface import MyInterface diff --git a/examples/marker_streaming.py b/examples/marker_streaming.py index 1d52244..4cccea0 100644 --- a/examples/marker_streaming.py +++ b/examples/marker_streaming.py @@ -17,6 +17,7 @@ If you want to display the markers in a 3D scatter plot, you can add a Scatter3D plot to the interface. You can pass the size and color of the marker via size and color argument, respectively. Please see the Scatter3D documentation for more information. Next, the data flow runs in a loop where the get_marker_set_data() function is used to retrieve the data from the interface. The data is then passed to the graph via the update() method through an array of (n_frame, n_markers, 3) where the plot parameters can be updated. """ + from time import sleep, time from custom_interface import MyInterface from biosiglive import LivePlot, PlotType, ViconClient diff --git a/examples/stream_data.py b/examples/stream_data.py index bf34a0c..e932824 100644 --- a/examples/stream_data.py +++ b/examples/stream_data.py @@ -16,6 +16,7 @@ at the save frequency specified in the start method. Please note that it is not yet possible to plot the data in real-time. """ + from custom_interface import MyInterface from biosiglive import ( ViconClient, diff --git a/examples/trigno_sdk.py b/examples/trigno_sdk.py new file mode 100644 index 0000000..8558530 --- /dev/null +++ b/examples/trigno_sdk.py @@ -0,0 +1,18 @@ +from biosiglive import TrignoSDKClient +from biosiglive import LivePlot, PlotType + +if __name__ == "__main__": + plot_curve = LivePlot( + name="curve", + rate=100, + plot_type=PlotType.Curve, + nb_subplots=1, + channel_names=["1"], + ) + plot_curve.init(plot_windows=10000, y_labels=["emg"]) + + client = TrignoSDKClient() + client.start_streaming() + while True: + data, timestamp = client.all_queue["avanti_emg"].get() + plot_curve.update(data[0:1, :]) diff --git a/pyproject.toml b/pyproject.toml index 92e8005..35cd05e 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,7 +5,7 @@ build-backend = "hatchling.build" [project] name = "biosiglive" -version = "2.0.2" +version = "2.0.4" authors = [ { name="Amedeo ceglia", email="amedeo.ceglia@umontreal.ca" }, ] diff --git a/tests/test_data_processing.py b/tests/test_data_processing.py index d491f7f..11071b3 100644 --- a/tests/test_data_processing.py +++ b/tests/test_data_processing.py @@ -1,3 +1,5 @@ +from numpy import lib +from numpy.compat import Path import pytest from biosiglive import ( RealTimeProcessing, @@ -76,7 +78,6 @@ def test_offline_processing(method): np.random.seed(50) data = np.random.rand(2, 4000) processing = OfflineProcessing(data_rate=2000, processing_window=1000) - if method == OfflineProcessingMethod.ProcessEmg: processed_data = processing.process_emg(data) np.testing.assert_almost_equal(processed_data[:, 0], [0.1021941, 0.0959892]) @@ -95,13 +96,15 @@ def test_offline_processing(method): def test_inverse_kinematics_methods(methods): parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) markers_data = np.ones((3, 16, 1)) * 0.1 - model_path = parent_dir + "/examples/model/Wu_Shoulder_Model_mod_wt_wrapp.bioMod" - + model_path = os.path.join(parent_dir, "examples", "model", "Wu_Shoulder_Model_mod_wt_wrapp.bioMod") + if os.name == "nt": # If the system is Windows, we need to replace the drive letter with an uppercase one + p = Path(model_path) + model_path = str(p).replace(f"{p.drive}", p.drive.upper(), 1) msk_function = MskFunctions(model=model_path) - q, q_dot = None, None + q, q_dot, qddot = None, None, None i = 0 while i != 15: - q, q_dot = msk_function.compute_inverse_kinematics(markers_data, method=methods) + q, q_dot, qddot = msk_function.compute_inverse_kinematics(markers_data, method=methods) i += 1 if methods == InverseKinematicsMethods.BiorbdLeastSquare: @@ -127,6 +130,7 @@ def test_inverse_kinematics_methods(methods): decimal=4, ) np.testing.assert_almost_equal(q_dot[:, 0], np.zeros((15,))) + np.testing.assert_almost_equal(qddot[:, 0], np.zeros((15,))) if methods == InverseKinematicsMethods.BiorbdKalman: np.testing.assert_almost_equal( @@ -171,14 +175,37 @@ def test_inverse_kinematics_methods(methods): ], decimal=4, ) + np.testing.assert_almost_equal( + qddot[:, 0], + [ + -0.086063, + -0.6521736, + 0.3024718, + 27.344666, + 3.5197842, + 1.6759323, + -41.2854412, + -29.9048829, + 6.6074331, + -8.3143703, + -4.7288301, + -0.6923808, + -0.4416418, + 4.6232351, + -0.161266, + ], + decimal=4, + ) def test_forward_kinematics(): parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) np.random.seed(50) q_data = np.random.rand(15, 1) - model_path = parent_dir + "/examples/model/Wu_Shoulder_Model_mod_wt_wrapp.bioMod" - + model_path = os.path.join(parent_dir, "examples", "model", "Wu_Shoulder_Model_mod_wt_wrapp.bioMod") + if os.name == "nt": # If the system is Windows, we need to replace the drive letter with an uppercase one + p = Path(model_path) + model_path = str(p).replace(f"{p.drive}", p.drive.upper(), 1) msk_function = MskFunctions(model=model_path) markers = None i = 0 diff --git a/tests/test_interfaces.py b/tests/test_interfaces.py index 0f457ba..40cd15a 100644 --- a/tests/test_interfaces.py +++ b/tests/test_interfaces.py @@ -1,3 +1,5 @@ +from pathlib import Path + import pytest from biosiglive import ( InterfaceType, @@ -102,6 +104,9 @@ def test_devices(device_type): def test_marker_set(): parent_dir = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) model_path = parent_dir + "/examples/model/Wu_Shoulder_Model_mod_wt_wrapp.bioMod" + if os.name == "nt": # If the system is Windows, we need to replace the drive letter with an uppercase one + p = Path(model_path) + model_path = str(p).replace(f"{p.drive}", p.drive.upper(), 1) marker_set = MarkerSet(nb_channels=16, name="my_marker_set", rate=100, system_rate=100) marker_set.data_window = 100 assert marker_set.sample == 100 / 100 @@ -116,6 +121,6 @@ def test_marker_set(): i += 1 raw_data = marker_set.raw_data - kin_data, _ = marker_set.get_kinematics(model_path=model_path, method=InverseKinematicsMethods.BiorbdKalman) + kin_data, _, _ = marker_set.get_kinematics(model_path=model_path, method=InverseKinematicsMethods.BiorbdKalman) assert raw_data.shape == (3, 16, 100) assert kin_data.shape == (15, 1) diff --git a/tests/test_io.py b/tests/test_io.py index 6979c0f..bd37a30 100644 --- a/tests/test_io.py +++ b/tests/test_io.py @@ -18,7 +18,7 @@ def test_save_and_load(rt): i = 0 while i != 50: if rt: - save(data, "test") + save(data, "test", add_data=True) i += 1 if not rt: save(data, "test") @@ -26,6 +26,7 @@ def test_save_and_load(rt): os.remove("test.bio") np.testing.assert_almost_equal(len(list(data_loaded.keys())), len(list(data.keys()))) np.testing.assert_almost_equal(data_loaded["data_np"].shape, (2, shape_np)) - np.testing.assert_almost_equal(len(data_loaded["data_int"]), shape) - np.testing.assert_almost_equal(len(data_loaded["data_list"]), shape) - np.testing.assert_almost_equal(len(data_loaded["data_dict"]), shape) + if rt: + np.testing.assert_almost_equal(len(data_loaded["data_int"]), shape) + np.testing.assert_almost_equal(len(data_loaded["data_list"]), shape * len(data["data_list"])) + np.testing.assert_almost_equal(len(data_loaded["data_dict"]), 1)