diff --git a/.gitignore b/.gitignore index 32b49c8..c773fda 100644 --- a/.gitignore +++ b/.gitignore @@ -84,3 +84,7 @@ __pycache__/ *.pyc figures/ .venv/ + +# Generated CAP-window caches & trained model weights (large, reproducible) +scripts/cap_pipeline/caches/ +*.pt diff --git a/articles/2022.10.05.510983v1.full.pdf b/articles/2022.10.05.510983v1.full.pdf new file mode 100644 index 0000000..ffbcee0 Binary files /dev/null and b/articles/2022.10.05.510983v1.full.pdf differ diff --git a/articles/2022_Resource_efficient_Neural_Network_Architectures_for_Classifying_Nerve_Cuff_Recordings.md b/articles/2022_Resource_efficient_Neural_Network_Architectures_for_Classifying_Nerve_Cuff_Recordings.md new file mode 100644 index 0000000..e527636 --- /dev/null +++ b/articles/2022_Resource_efficient_Neural_Network_Architectures_for_Classifying_Nerve_Cuff_Recordings.md @@ -0,0 +1,336 @@ +# Content from 2022.10.05.510983v1.full.pdf + +bioRxiv preprint doi: https://doi.org/10.1101/2022.10.05.510983; this version posted October 7, 2022. The copyright holder for this preprint (which was not certified by peer review) is the author/funder. All rights reserved. No reuse allowed without permission. + +# Resource-efficient Neural Network Architectures for Classifying Nerve Cuff Recordings on Implantable Devices + +Authors: Y. E. Hwang^1,2^, R. Genov^3^, J. Zariffa^1,2,3,4^; + +^1^Institute of Biomedical Engineering, University of Toronto, Toronto, Canada; + +^2^KITE Research Institute, Toronto Rehabilitation Institute – University Health Network, Toronto, Canada + +^3^Edward S. Rogers Sr. Department of Electrical & Computer Eng., University of Toronto, Toronto, Canada + +^4^Rehabilitation Sciences Institute, University of Toronto, Toronto, Canada + +Corresponding author: jose.zariffa@utoronto.ca + +Disclosures: + +Y. E. Hwang: None. R. Genov: None. J. Zariffa: None. + +Abstract: + +**Background:** Closed-loop control of functional electrical stimulation involves using recorded nerve signals to make decisions regarding nerve stimulation in real-time. Surgically implanted devices that can implement this strategy have significant potential to restore natural movement after paralysis. Previous work demonstrated the use of convolutional neural networks (CNNs) to discriminate between activity from different neural pathways recorded by a high-density multi-contact nerve cuff electrode. Despite state-of-the-art performance, that approach required too much data storage, power and computation time for a practical implementation on surgically implanted hardware. + +**Objective:** To reduce resource utilization for an implantable implementation, with a minimal performance loss for CNNs that can discriminate between neural pathways in multi-contact nerve cuff electrode recordings. + +**Methods:** Neural network (NN) architectures were evaluated on a dataset of rat sciatic nerve recordings previously collected using 56-channel (7 x 8) spiral nerve cuff electrodes to capture spatiotemporal neural activity patterns. The NNs were trained to classify individual, natural compound action potentials (nCAPs) elicited by sensory stimuli. Three architecture types were explored: the previously reported ESCAPE-NET, a fully convolutional network, and a recurrent neural network. Variations of each architecture yielded NNs with a range in the number of weights and required floating-point operations (FLOPs). Each NN was evaluated based on F1-score and resource requirements. + +**Results:** NNs were identified that, when compared to ESCAPE-NET, required 1,132-1,787x fewer weights, 389-995x less memory, and 6-11,073x fewer FLOPs, while maintaining macro F1-scores of 0.70-0.71 compared to a baseline of 0.75. Memory requirements range from 22.69 KB to 58.11 KB, falling within the range of on-chip memory sizes from several published deep learning accelerators fabricated in 65nm ASIC technology. + +**Conclusion:** Reduced versions of ESCAPE-NET require significantly fewer resources without significant accuracy loss, thus can be more easily incorporated into a surgically implantable device that performs closed-loop real-time responsive neural stimulation. + +--- + +## 1 Introduction + +Afferent signals contain proprioceptive and sensory information, which can be used to determine joint position, muscle stretch and tension, and presence of external stimuli. The central nervous system (CNS) uses such information to guide limb movement and reflex responses [1]; in the case of paralysis due to spinal cord injury, supraspinal control of limb movements may be disrupted, preventing the muscles from receiving motor commands from the brain. + +Open-loop functional electrical stimulation (FES) can be delivered via peripheral nerves to restore motor function in paralyzed limbs, but without the sophistication of motor commands generated by the brain. This results in coarse limb movement lacking the nuance of unimpaired motor control, the inability to compensate for unexpected perturbations or errors, and inefficient stimulation delivery that causes muscle fatigue [2]–[4]. Closed-loop FES aims to improve stimulation delivery by using afferent signals as feedback to regulate the stimulation delivered. + +The information-carrying afferent signals can be recorded using nerve interfacing electrodes, after which data processing can be applied to extract and interpret the proprioceptive and sensory information. Several studies have demonstrated using such information as feedback control to achieve closed-loop FES in animal models. Bruns et al. [5] demonstrated joint angle decoding by applying linear regression on dorsal root ganglion (DRG) signals to regulate FES of various leg muscles to produce a regular gait cycle. Holinski et al. [6] demonstrated ground reaction force prediction during stepping by multivariate linear equations on DRG signals to regulate intraspinal microstimulation to control hind leg motion. Song et al. [7] demonstrated ankle angle estimation by fast independent component analysis and dynamically driven recurrent neural network on sciatic nerve signals to regulate FES to control ankle motion. Closed-loop FES has also been demonstrated in human trials. Wendelken et al. [8] demonstrated closed-loop control of a virtual prosthetic hand using intraneural electrodes on the human median nerve decoded by a Kalman filter. Hansen et al. [9] demonstrated closed-loop foot drop correction using human sural nerve recordings discriminated by an artificial neural network. Inmann et al. [10] demonstrated a closed-loop hand grasp system using human median nerve recordings processed by rectification and bin-integration, high- and low-pass filtering, and delay to control hand grasp to react to and stop slipping objects. + +Peripheral nerve signals may be recorded by surgically implanted devices of varying levels of invasiveness [11]. Extraneural nerve cuffs are advantageous as they do not penetrate the nerve and are typically preferable for clinical use due to long-term stability [12], [13]. However, naturally-evoked compound action potentials (nCAP) exhibit low signal-to-noise ratio (SNR). Thus, achieving recording selectivity (i.e. the ability to determine which neural pathway produced the measured signals) using extraneural devices is challenging [11]. Improving the selectivity of peripheral nerve recordings is necessary to provide functionally specific feedback signals that can improve the performance of closed-loop FES systems. + +Koh et al. recently demonstrated that selective recording of individual nCAPs from extraneural nerve cuff electrodes was possible. Deep learning was used to classify sciatic nerve recordings obtained by multi-contact nerve cuffs with state-of-the-art selectivity [14]. They used multi-contact nerve cuffs with multiple rings of contacts along the length of the nerve and multiple contacts per ring around the diameter of the nerve; this structure captures velocity and spatial variability, respectively, between different neural pathways. The novel method, named Extraneural Spatiotemporal CAPs Extraction Network (ESCAPE-NET), detects individual nCAPs then represents them as a 2-dimensional spatiotemporal signature. ESCAPE-NET uses a convolutional neural network (CNN) to classify spatiotemporal signatures. Classifying individual nCAPs is particularly advantageous as it provides high temporal resolution and could potentially handle scenarios where multiple neural pathways are active, which complicates classification when using longer time windows. + +To be successfully translated to humans, closed-loop FES technology must be implemented on surgically implantable devices. A surgically implantable device would exhibit minimal interference with a user's daily routine and would draw minimal attention to the person's disability due to being externally invisible [15] [16]. However, neural networks (NN) traditionally rely on computational complexity to achieve high classification accuracy for non-linearly separable datasets. Computational complexity translates to longer computation time and size of hardware in terms of either memory, accelerators, or power consumption, which hinders deployment onto surgically implantable devices with strict hardware constraints. Recent efforts on reducing a NN's resource intensiveness with minimal performance degradation include MobileNet [17] [18] [19], EfficientNet [20] [21], and Once-For-All [22] networks. Specifically, focus has been placed on reducing the number of parameters or weights, which are crucial to the NN's capacity to learn and detect structural features in the input data. Springenberg et al. developed a fully convolutional network that outperformed more complicated state-of-the-art networks, achieving large parameter savings that doubles as a form of regularization [23]. Another benefit of reducing the computational burden of a NN is the reduction of latency (computational time for inference) required for a NN to process its inputs. + +To combine the proven benefits of using NNs to discriminate between neural signals and the goal of real-time performance on implantable hardware, there is a need to explore reduced NNs that perform with minimal accuracy loss and are suitable for real-time performance on implantable hardware. This study aimed to reduce resource consumption with minimal performance accuracy loss for machine learning techniques that can distinguish between neural pathways in multi-contact nerve cuff electrode recordings. + +--- + +## 2 Methods + +### 2.1 Dataset description + +The comparison of the neural network architectures made use of an existing dataset consisting of spatiotemporal signatures extracted from detected nCAPs in the rat sciatic nerve. More details regarding the dataset and how it was collected can be found in [24]. Figure 1 illustrates detected nCAPs and their corresponding spatiotemporal signatures. There are three classes of spatiotemporal signatures in this dataset, labeled based on the type of stimulus that was being applied at the time each nCAP was recorded: ankle dorsiflexion, ankle plantarflexion, and cutaneous stimulus to the heel using a Von Frey monofilament. Data was collected from 9 Long-Evan rats, with a separate neural network trained for each rat. 3-fold cross-validation was used in each case. + +Electroneurographic (ENG) signals were recorded at 30 kHz using the 56 contacts of the nerve cuff electrode, which were arranged as 7 rings of 8 contacts. Each spatiotemporal signature contained a 3.33 ms segment of the recorded signal, centered around the detected peak in the middle ring of contacts. Each spatiotemporal signature therefore had dimensions of 56 channels x 100 time samples. Two representations of the spatiotemporal signature are possible: the first groups channels by nerve circumference (i.e. every 8 consecutive rows correspond to a ring of contacts) and the second groups channels by nerve length (i.e. every 7 consecutive rows correspond to a line of contacts along the nerve). + +The number of nCAPs per category was 3,898 ± 2,175 for dorsiflexion, 3,454 ± 2,433 for plantarflexion, and 3,991 ± 2,145 for heel prick. Training sets for each rat were augmented to include a total of 10,000 training examples of spatiotemporal signatures per class. Augmented nCAPs were generated by averaging all nCAPs associated with a label then adding randomly sampled noise to create new noisy nCAPs [14]. Testing sets included only detected nCAPs (no augmented nCAPs). The validation set was comprised of a randomly sampled 10% of the training set. + +--- + +*Figure 1. Example of the ENG signals used for classification. For clarity of visualization, the depicted signals are averaged across several thousands of nCAPs. Inputs to the NNs consisted of single nCAP recordings, which would contain higher levels of noise. (a) Processed signals of a detected nCAP measured on all contacts of the nerve cuff electrode. The signals have been filtered, tripole referenced and normalized between -1 and 1 before being averaged. (b) Spatiotemporal signature representation of the ENG of the signals presented in (a); contacts are numbered based on the green numbers shown in (a) to the right of each plotted nCAP. (c) Spatiotemporal signature representation of the ENG of the signals presented in (a); contacts are numbered based on the orange numbers in (a) on top of each plotted nCAP.* + +--- + +### 2.2 Evaluation of neural networks + +Several NN architectures were explored for the purpose of classifying spatiotemporal signatures and evaluated based on the amount of required resources and classification ability. The goal was to determine whether or not there exists a NN architecture that uses an amount of resources compatible with eventual application-specific integrated circuit (ASIC) translation while demonstrating competitive classification performance. + +For the purposes of this study, we define resources as a network's required memory and floating-point operations (FLOPs), which are directly influenced by the network's architecture, the number of weights in a network, and the precision of the numeric representation of the weights and intermediate values. We define a network's classification ability based on accuracy and macro F1-score of the network's classification of a labelled test set. For each datum in the test set, the label with the highest softmax output of the neural network is the predicted label. + +Number of parameters was measured using the built-in Keras function count_params. Amount of memory was measured using the size of Tensorflow Lite (TFLite) converted models. FLOPs were measured using the Tensorflow Profiler. When determining classification accuracy and F1-score, the label with the highest predicted probability is compared to the ground truth label for equivalence. + +### 2.3 Implementation and training of neural networks + +All neural network models were designed and trained using Tensorflow Keras libraries. Training was performed over 1000 epochs or whenever the validation loss function did not decrease for over 15 epochs. The optimization algorithm used for network training was stochastic gradient descent (SGD) with a momentum value of 0.9. The loss function used was categorical cross entropy (CCE) loss. All networks ended with a dense layer responsible for classification using the softmax activation function. + +### 2.4 Neural network architectures + +We refer to an architecture as a specific type of network. The first architecture was ESCAPE-NET, the CNN previously published by Koh et al. [14]. The second architecture was a Fully Convolutional Network (FCN), as inspired by Lin et al. [25], using smaller filter sizes and deeper networks as inspired by multiple recent successful neural networks [26]. Fully convolutional networks replace traditional dense layers with convolutional layers that perform the same functionality and have been shown to reduce overfitting and promote parameter efficiency. The third architecture was a Recurrent Neural Network (RNN), which was explored to exploit the time series nature of the dataset in question. + +Within a single architecture, structural details were varied to explore networks with different trade-offs between amount of resources and performance. Examples include the network width (i.e. the number of filters per layer), the number of layers, filter sizes, regularization, or types of layers within the same family (e.g. Long-Short Term Memory and Gated Recurrent Unit layers for the RNN family). + +Each structural detail was explored sequentially. After each structural detail was explored, the network that yielded the fewest number of weights while maintaining performance higher than 0.7 macro F1-score was chosen. Each structural detail remained fixed based on this requirement for the remainder of the exploration. + +#### 2.4.1 ESCAPE-NET Mini + +Introduced by Koh et al. [14], ESCAPE-NET receives two inputs, corresponding to two different representations of the spatiotemporal signature. The two representations differ in the ordering of the contacts. The "spatial emphasis" representation ensures that rows of pixels are locally related based on contact placement around the diameter of the nerve, whereas the "temporal emphasis" representation ensures that rows of pixels are locally related based on contact placement along the length of the nerve (Figure 1). The former exploits the spatial variation of different neural pathways, and the latter exploits the velocity variation of different neural pathways. Reduced versions of ESCAPE-NET, hereon noted as ESCAPE-NET Mini, are explored by varying the following structural details detailed in this section. + +For each of the two inputs of the network, there are multiple convolutional layers using the ReLU activation function in sequence, separated by 2x2 max pooling layers for data downsampling. To initialize the weights of the first and second convolutional layers, they were temporarily connected to dense layers and trained for 25 epochs. The pair of convolutional layer stacks were concatenated after the third convolutional layer into a dense layer using the ReLU activation function. Figure 2a illustrates a sample ESCAPE-NET network explored in this work. + +An initial sweep of network widths was performed by varying the convolutional layer width and dense layer width between 2-64, increasing by a factor of 2, keeping the two values equal throughout this sweep. A specific value was selected after this initial sweep using the network with the fewest weights performing above a macro F1-score of 0.7. During this sweep, the number of convolutional layers was set to 3, the filter sizes were set to 8x8 for conv layer 1, 4x4 for conv layer 2, 2x2 for conv layer 3, and the dropout value was set at 0. + +To investigate the effect of unequal convolutional and dense layer widths, a secondary sweep of network widths was performed on a smaller search space centered around the chosen value from the initial sweep. For example, if a width of 8 was chosen after the initial sweep, convolutional and dense widths were explored for all combinations of 4, 8 and 16 (i.e. 9 combinations in total). Specific values were selected for convolutional and dense layer widths after this secondary sweep. + +The number of layers was varied between 2-4. Filter sizes were varied between 8x8 for conv layer 1, 4x4 for conv layer 2, 2x2 for conv layer 3 and 9x9 for conv layer 1, 5x5 for conv layer 2, 3x3 for conv layer 3. The dropout value was varied between 0 and 0.5 to represent the absence and presence of a dropout layer, respectively. The dropout value of 0.5 was used due to its reported efficacy [27]. + +#### 2.4.2 Fully Convolutional Network (FCN) + +The second explored architecture was an FCN using smaller filter sizes and a deeper network. Like ESCAPE-NET, it receives two inputs for the different representations of the spatiotemporal signature. For each of the two inputs of the network, there are multiple sets of convolutional layers using the ReLU activation function and filter size of 3x3 in sequence, separated by 2x2 max pooling layers for data downsampling. The second and third sets of convolutional layers were followed by a 1x1 convolutional layer to allow for more complex function modelling. The pair of convolutional layer stacks were concatenated after the third set of convolutional layers into a 1x1 convolutional layer using the ReLU activation function in place of the dense layer in ESCAPE-NET, and a 1x1 convolutional layer and Global Average Pooling layer in place of the dense layer responsible for classification. Figure 2b illustrates a sample FCN network explored in this work. + +An initial sweep of network widths was performed by varying the number of filters of the first set of convolutional layers (i.e. base convolutional layer width) between 2-64, increasing by a factor of 2, and varying the number of filters of the 1x1 layer after concatenation (i.e. 1x1 convolutional layer width) between 2-64, increasing by a factor of 2, keeping the two values equal throughout this sweep. A specific value was selected for each width after these initial sweeps based on its resource/performance trade-off. During this initial sweep, the number of sets of layers was set to 3, and the channel width multiplier was set to 2. The channel width multiplier determines how layer width increases from one set of convolutional layers to the next. + +To investigate the effect of unequal base convolutional layer and 1x1 convolutional layer widths, all combinations of the selected width, half the selected width and double the selected width were evaluated (i.e. 9 networks in total). A specific value was selected for each width after this exploration. + +Structural details that were varied for this architecture were the number of sets of layers and channel width multiplier. The choice to vary these particular structural details was based on CNN advancements that found success using deeper networks with smaller filter sizes [28], [29], increasing width of layers in deeper layers [26], and 1x1 convolutions in the place of dense layers [25]. + +The number of sets of convolutional layers was varied between 2-4. The channel width multiplier value was varied between 1-3. + +#### 2.4.3 Recurrent Neural Network (RNN) + +The third explored architecture was an RNN consisting of multiple recurrent layers using the tanh activation function connected in sequence, followed by a dense layer using the ReLU activation function. Unlike the previous architectures, there is only one input to this network, as each contact of the nerve cuff is treated as a single channel of the data, eliminating the benefit of having localized spatial or temporal information encoded in the spatiotemporal signature. Figure 2c illustrates a sample RNN network explored in this work. + +An initial sweep of network widths was performed by varying the number of units of the recurrent layer (i.e. recurrent layer width) and the dense layer width between 2-64, increasing by a factor of 2, keeping the two values equal throughout this sweep. A specific value was selected for each width after these initial sweeps based on its resource/performance trade-off. During these sweeps, the number of layers was set as 3, and the layer type was set as GRU. + +To investigate the effect of unequal recurrent layer and dense layer widths, all combinations of the selected width, half the selected width and double the selected width were evaluated (i.e. 9 networks in total). A specific value was selected for each width after this exploration. + +Structural details that were varied for this architecture were the number of layers, number of neurons per layer, and layer type. The number of layers of the network was varied from 2-5 given a fixed number of neurons per layer. The number of layers was fixed after this exploration. Layer types that were explored consisted of Long-Short Term Memory (LSTM), Gated Recurrent Unit (GRU), Bidirectional LSTM, and Bidirectional GRU. The last layer was varied between the recurrent layer and an attention layer. + +### 2.5 Lower weight precision and memory requirement definition + +Lower precision numerical representation of weights was explored due to its direct relevance to minimizing hardware requirements. Using fewer bits to represent a numerical value reduces the amount of memory required to store values and the complexity of multiply and accumulate operations, but potentially sacrifices accuracy due to information loss. By default, Tensorflow models use 32-bit floating point numbers to store weights. We explored using 8-bit integer weights using the TensorFlow Lite default optimization from the TensorFlow Model Optimization Toolkit. A TFLITE file is produced that contains the neural network model's execution graph. We used the size of a TFLITE file to define a network's memory requirement. + +--- + +*Figure 2. Network design for the novel architectures explored: (a) ESCAPE-NET Mini architecture, (b) ESCAPE-FCN architecture, and (c) ESCAPE-RNN architecture. Structural detail values used in these illustrations are based on the highlighted networks presented in the Results section.* + +--- + +## 3 Results + +We present the results of the network exploration performed by varying structural detail values for each of the three architectures. Each network's classification accuracy, macro F1-score, number of weights and FLOPs are presented. For each architecture, a single network is highlighted and analyzed in more detail. + +Figure 3 displays the calculated number of weights for all the explored networks for each of the three architectures measured against their achieved macro F1-score. Figure 4 displays the calculated number of FLOPs for the same networks against their achieved macro F1-score. A single datapoint represents the average performance of a network across all 9 subjects and 3 folds. A highlighted network, as indicated by the markers with black borders, was selected for each of the three architectures based on the architecture requiring the fewest weights that meets our performance requirement of 0.70 macro F1-score. + +As expected, the performance of the networks within each architecture trended downwards as the number of resources decreased. The majority of the explored networks performed within 0.10 macro F1-score of the baseline network from [14]. + +--- + +*Figure 3. Resource/performance trade-off plot for every explored network configuration for each of the three architectures, analyzing the macro F1-score for performance versus the number of weights for resources. Each data point represents the average performance of a network design trained with 27 datasets (9 subjects, 3-fold cross validation). The dotted line represents the minimum performance constraint used in this study. The points with black borders represent the highlighted network configurations selected for each of the 3 architectures as well as the original ESCAPE-NET architecture.* + +--- + +*Figure 4. Resource/performance trade-off plot for every explored network configuration for each of the 3 architectures, analyzing the macro F1-score for performance versus the number of FLOPs for resources.* + +### 3.1 ESCAPE-NET Mini architecture + +The initial sweep of layer widths was performed by varying the convolutional and dense layer width values between 2, 4, 8, 16, 32, and 64, keeping the two values equal. The convolutional and dense layer widths were chosen to be 8 after this initial sweep. The secondary sweep of layer widths was performed by varying the convolutional and dense layer width values between 4, 8 and 16, exploring all combinations. The convolutional layer width was chosen to be 4 and the dense layer width was chosen to be 8 after this secondary sweep. + +The number of convolutional layers was chosen to be 4 after exploring networks with 2 layers, 3 layers and 4 layers. The first layer uses 8x8 filters, the second layer uses 4x4 filters, the third layer uses 2x2 filters, and the fourth layer uses 1x1 filters. Dropout was not used after the dense layer. + +### 3.2 ESCAPE-FCN architecture + +The initial sweep of layer widths was performed by varying the base convolutional and 1x1 convolutional layer width values between 2, 4, 8, 16, 32, and 64, keeping the two values equal. The base convolutional and 1x1 convolutional layer widths were chosen to be 16 after this initial sweep. The secondary sweep of layer widths was performed by varying the base convolutional and 1x1 convolutional layer width values between 8, 16 and 32, exploring all combinations. The base convolutional layer width was chosen to be 16 and the 1x1 convolutional layer width was chosen to be 8 after this secondary sweep. + +The number of sets of convolutional layers was chosen to be 2 after exploring networks with 2, 3 and 4 sets of layers. The channel width multiplier value was chosen to be 2. + +### 3.3 ESCAPE-RNN architecture + +The initial sweep of layer widths was performed by varying the recurrent and dense layer width values between 2, 4, 8, 16, 32, and 64, keeping the two values equal. The recurrent and dense layer widths were chosen to be 32 after this initial sweep. The secondary sweep of layer widths was performed by varying the recurrent and dense layer width values between 16, 32 and 64, exploring all combinations. The recurrent layer width was chosen to be 32 and the dense layer width was chosen to be 16 after this secondary sweep. + +The number of layers was chosen to be 3. The 1^st^ and 2^nd^ layers were chosen to be LSTM, and the 3^rd^ layer was an attention layer. + +--- + +*Figure 5. Comparisons between each highlighted network configuration with the baseline ESCAPE-NET network from [14]. (a) Comparison of number of weights and macro F1-score across all 4 networks. (b) Comparison of memory footprint and macro F1-score across all 4 networks. The horizontal dotted lines represent the memory requirements from cited works [30]–[33] respectively. (c) Comparison of number of FLOPs and macro F1-score across all 4 networks.* + +--- + +### 3.4 Resource/Performance trade-off analysis + +To analyze the resource/performance plots presented in the Results section, networks nearer to the top left corner are preferable. However, this study is not strictly an optimization problem for minimal resources and maximum performance. Alternative ways of interpreting the plots are by defining a resource constraint given a specific hardware platform then selecting the network that achieves the best performance within this constraint or defining a minimum performance requirement and selecting the network that uses the fewest resources above this threshold. + +To select a highlighted network for each architecture as presented in Figure 3, we use a performance constraint of a macro F1-score threshold of 0.7. In Figure 5b, we use resource constraints specified by published deep learning accelerator implementations. These resource and performance constraints are specific to this study. Different constraints may be defined based on different applications and platforms. Within this selection of networks, there are many alternatives to the highlighted networks that may be suitable for practical use. + +### 3.5 Highlighted networks + +Table 1 summarizes the evaluated metrics of the three highlighted networks. These are the networks for each architecture that require the fewest weights while performing with macro F1-score of above 0.70. Of the three selected networks, ESCAPE-NET Mini uses the fewest weights and therefore memory, and ESCAPE-RNN uses the fewest FLOPs. ESCAPE-NET Mini can achieve better performance at higher amounts of resources, whereas the other two architectures plateau. + +Table 2 compares the performance and memory for the three highlighted networks with and without lower precision weights. Every network performs with approximately equal accuracy and macro F1-score after conversion to lower precision weights. All reported instances of memory footprint in this document outside of Table 2 uses the TF optimized model with 8-bit integer weights. + +*Table 1. Summary of the performance and resource requirements of the baseline network and the 3 highlighted networks.* + +| Architecture | Classification Accuracy (%) | Macro F1-Score | Number of Weights | Memory (KB) | FLOPs | +|---|---|---|---|---|---| +| ESCAPE-NET [14] | 80.76 ± 10.13 | 0.75 ± 0.11 | 23,112,067 | 22,585 | 529,405,202 | +| ESCAPE-NET Mini | 77.22 ± 11.64 | 0.70 ± 0.13 | 12,931 | 22.69 | 7,405,290 | +| ESCAPE-FCN | 77.71 ± 10.79 | 0.71 ± 0.12 | 15,075 | 27.32 | 83,549,215 | +| ESCAPE-RNN | 78.41 ± 10.01 | 0.71 ± 0.12 | 20,423 | 58.11 | 47,810 | + +--- + +*Table 2. Performance and memory comparison between higher- and lower-precision weights.* + +**32-bit floating point weights** + +| Architecture | Classification Accuracy (%) | Macro F1-Score | Memory (KB) | +|---|---|---|---| +| ESCAPE-NET [14] | 80.77 ± 10.13 | 0.75 ± 0.11 | 90,287 | +| ESCAPE-NET Mini | 77.22 ± 11.64 | 0.70 ± 0.13 | 57.65 | +| ESCAPE-FCN | 77.70 ± 10.78 | 0.71 ± 0.12 | 65.70 | +| ESCAPE-RNN | 78.41 ± 10.01 | 0.71 ± 0.12 | 113.05 | + +**8-bit integer weights** + +| Architecture | Classification Accuracy (%) | Macro F1-Score | Memory (KB) | +|---|---|---|---| +| ESCAPE-NET [14] | 80.76 ± 10.13 | 0.75 ± 0.11 | 22,585 | +| ESCAPE-NET Mini | 77.22 ± 11.64 | 0.70 ± 0.13 | 22.69 | +| ESCAPE-FCN | 77.71 ± 10.79 | 0.71 ± 0.12 | 27.32 | +| ESCAPE-RNN | 78.44 ± 10.01 | 0.71 ± 0.12 | 58.11 | + +--- + +## 4 Discussion + +This study explored 3 different NN architectures, varying structural details for each to explore the trade-off between performance and accuracy for networks with reduced resource requirements compared to the baseline network from [14]. We found versions of each architecture with significantly reduced resource requirements that maintained macro F1-score of above 0.70. In particular, the networks' memory requirements fell by an order of magnitude from 10s of megabytes to 10s of kilobytes. + +### 4.1 Comparison between highlighted networks + +With structural detail tuning, a version of each architecture could achieve comparable trade-offs between the number of weights and macro F1-score. ESCAPE-NET Mini performs slightly better when given more weights, and ESCAPE-RNN performs slightly better when given fewer weights. These networks may be better choices if the hardware constraints in a practical setting are loosened or tightened, respectively. Hardware constraints may be loosened due to implementation of the neural network on a wearable device or cloud storage in communication with the implant. + +However, when considering FLOPs, ESCAPE-RNN networks are clearly advantageous, and ESCAPE-FCN networks are clearly disadvantageous. If processing speed is a bottleneck, an ESCAPE-RNN network may be preferred. + +### 4.2 Defining resource constraints + +Whether a particular neural network can be implemented onto a particular target device with real-time performance is a complex question that is difficult to determine prior to actual implementation. However, as this work was intended to facilitate eventual hardware implementation, we are interested in defining an informed estimate of potential memory and latency constraints and analyzing the neural networks we presented that fall within these limitations. Based on specifications of published deep learning implementations on ASIC technology and the memory requirements of our neural networks, we argue that our neural networks can be implemented on 65nm technology. + +Moons et al. achieved a low-power high-precision embedded processor designed for large scale CNNs implemented in 40nm technology and equipped with 148 KB of on-chip memory [32]. Both Qiao et al. and Yue et al. achieved Computing-in-Memory-based processors for NNs implemented in 65nm technology equipped with 73 KB [33] and 164 KB [31] of on-chip SRAM, respectively. Chen et al. achieved an energy-efficient accelerator for Deep CNNs implemented in 65nm technology and equipped with 181.5 KB of on-chip SRAM [30]. As illustrated in Figure 5b, each of this work's proposed networks' memory requirements fall within the amount of on-chip SRAM available on these neural network hardware implementations. + +Power consumption is a related metric to memory use and impacts the usability of implantable devices. Power consumption of neural networks depends on characteristics of the hardware platform on which they are implemented and the energy required to read weights from memory and perform each operation [34]. This study's focus on decreasing the number of weights and FLOPs will have direct impact on decreasing the power consumption of the networks once implemented on a hardware platform. + +### 4.3 Defining performance constraints + +Defining a particular performance requirement for this technology in a clinical setting will be highly subject and application based. A macro F1-score threshold of 0.7 was defined for this study based on the baseline macro F1-score of 0.747. This value is based on an average over 9 subjects, but this technology's intended clinical use will be based on performance for individual subjects. Figure 6 shows the resource/performance trade-off of the explored networks using the best and worst performing subjects rather than by averaging all subjects. The worst performing subjects exhibited far lower macro F1-scores with values nearer to 0.40. For these subjects, a larger network may be required for closed-loop FES to succeed. Conversely, the best performing subjects exhibited higher macro F1-scores with values nearer to 0.90. For these subjects, a smaller network may be sufficient for closed-loop FES to succeed. Given this variance in potential ideal networks for each subject, implementation of the neural network on a wearable device or cloud storage may be more practical due to less rigid constraints on network size. + +--- + +*Figure 6. Resource/performance trade-off plot for every explored network configuration for each of the 3 architectures. Each data point represents either the best or worst performing subject out of the 9 subjects used, using the average performance of the networks trained during 3-fold cross-validation. The dotted line represents the performance constraint used in this study.* + +--- + +### 4.4 Network convergence inconsistencies + +Due to the stochastic nature of neural network training, sometimes a network failed to converge for certain subjects. The 3 highlighted networks in this study converged for all subjects and folds. In this study, a network's final performance achieved after either the maximum number of training epochs or the early stopping criterion was met was incorporated for analysis regardless of whether or not the network converged. + +### 4.5 Future work + +This demonstration used the rat sciatic nerve due to its simplicity for an initial proof of concept. This method should be applied to more complex nerves, such as the human median nerve, to expand the potential reach of such technology. This demonstration was based on acute studies as a proof of concept for the proposed technology. Future work should progress towards chronic studies as this technology is intended for eventual chronic use. + +This study has demonstrated that more efficient machine learning techniques are effective at neural pathway discrimination. Future work should focus on implementing these neural networks on ASIC technology that are suitable for surgical implantation. The problem definition addressed in this work was whether NNs with reduced resources could perform with minimal performance loss compared to the original architecture. Future work should focus on performance improvements given resource constraints, potentially using a cost function to ensure an optimal solution is reached. + +--- + +## 5 Conclusion + +This study demonstrates that it is possible to achieve resource-efficient neural networks that perform with minimal accuracy loss compared to ESCAPE-NET and that are suitable for implementation on resource constrained implantable hardware. These efficient networks require a fraction of the resources spent by ESCAPE-NET, requiring 1,132-1,787x fewer weights, 389-995x less memory, and 6-11,073x fewer FLOPs, while maintaining macro F1-scores of over 0.70. The efficient networks have memory requirements within the constraints set by several published deep learning accelerators in 65nm ASIC fabrication technology, which is the targeted platform on which these efficient networks are intended to be implemented. + +The ability of efficient neural networks to discriminate between different types of peripheral nerve signals using a single extraneural device is a crucial step towards implementing closed-loop FES technology on implantable devices. Sophisticated closed-loop FES technology can restore nuanced natural control of paralyzed limbs, and implementation on implantable devices can ensure that the developed assistive device is not visible and does not affect the patient's quality of life. Beyond the impact on people with paralyzed limbs and spinal cord injury, this selective peripheral nerve recording technology has applications in the control of prosthetic limbs and neuromodulation interventions for diseases such as chronic pain, diabetes, or incontinence. + +--- + +## 6 Bibliography + +[1] U. Proske and S. Gandevia, "The proprioceptive senses: their roles in signaling body shape, body position and movement, and muscle force.," Physiol. Rev., 2012, doi: 10.1152/physrev.00048.2011. + +[2] P. H. Peckham and J. S. Knutson, "Functional Electrical Stimulation for Neuromuscular Applications," Annu. Rev. Biomed. Eng., vol. 7, no. 1, pp. 327–360, 2005, doi: 10.1146/annurev.bioeng.6.040803.140103. + +[3] T. S. Barss et al., "Utilizing Physiological Principles of Motor Unit Recruitment to Reduce Fatigability of Electrically-Evoked Contractions: A Narrative Review," Arch. Phys. Med. Rehabil., vol. 99, no. 4, pp. 779–791, Apr. 2018, doi: 10.1016/j.apmr.2017.08.478. + +[4] K. T. Gelenitis, B. M. Sanner, R. J. Triolo, and D. J. Tyler, "Selective Nerve Cuff Stimulation Strategies for Prolonging Muscle Output," IEEE Trans. Biomed. Eng., vol. 67, no. 5, pp. 1397–1408, May 2020, doi: 10.1109/TBME.2019.2937061. + +[5] T. M. Bruns, J. B. Wagenaar, M. J. Bauman, R. A. Gaunt, and D. J. Weber, "Real-time control of hind limb functional electrical stimulation using feedback from dorsal root ganglia recordings," J. Neural Eng., vol. 10, no. 2, p. 026020, Apr. 2013, doi: 10.1088/1741-2560/10/2/026020. + +[6] B. J. Holinski, D. G. Everaert, V. K. Mushahwar, and R. B. Stein, "Real-time control of walking using recordings from dorsal root ganglia," J. Neural Eng., vol. 10, no. 5, p. 056008, Oct. 2013, doi: 10.1088/1741-2560/10/5/056008. + +[7] K.-I. Song, S. E. Park, D. Hwang, and I. Youn, "Compact Neural Interface Using a Single Multichannel Cuff Electrode for a Functional Neuromuscular Stimulation System," Ann. Biomed. Eng., vol. 47, no. 3, pp. 754–766, Mar. 2019, doi: 10.1007/s10439-018-02181-1. + +[8] S. Wendelken et al., "Restoration of motor control and proprioceptive and cutaneous sensation in humans with prior upper-limb amputation via multiple Utah Slanted Electrode Arrays (USEAs) implanted in residual peripheral arm nerves.," J. Neuroengineering Rehabil., vol. 14, no. 1, Art. no. 1, Nov. 2017, doi: 10.1186/s12984-017-0320-4. + +[9] M. Hansen, M. Haugland, T. Sinkjaer, and N. Donaldson, "Real Time Foot Drop Correction using Machine Learning and Natural Sensors," Neuromodulation J. Int. Neuromodulation Soc., vol. 5, pp. 41–53, Jan. 2002, doi: 10.1046/j.1525-1403.2002._2008.x. + +[10] A. Inmann and M. Haugland, "Regulation of FES-induced grasp force based on cutaneous nerve signals: Experiments and modelling," Med. Eng. Phys., vol. 34, no. 1, Art. no. 1, Jan. 2012, doi: 10.1016/j.medengphy.2011.06.016. + +[11] C. E. Larson and E. Meng, "A review for the peripheral nerve interface designer," J. Neurosci. Methods, vol. 332, p. 108523, Feb. 2020, doi: 10.1016/j.jneumeth.2019.108523. + +[12] X. Navarro, T. B. Krueger, N. Lago, S. Micera, T. Stieglitz, and P. Dario, "A critical review of interfaces with the peripheral nervous system for the control of neuroprostheses and hybrid bionic systems," J. Peripher. Nerv. Syst. JPNS, vol. 10, no. 3, pp. 229–258, Sep. 2005, doi: 10.1111/j.1085-9489.2005.10303.x. + +[13] M. Ortiz-Catalan, R. Brånemark, B. Håkansson, and J. Delbeke, "On the viability of implantable electrodes for the natural control of artificial limbs: review and discussion," Biomed. Eng. Online, vol. 11, p. 33, Jun. 2012, doi: 10.1186/1475-925X-11-33. + +[14] R. G. L. Koh, M. Balas, A. I. Nachman, and J. Zariffa, "Selective peripheral nerve recordings from nerve cuff electrodes using convolutional neural networks.," J. Neural Eng., vol. 17, no. 1, p. 016042, Jan. 2020, doi: 10.1088/1741-2552/ab4ac4. + +[15] D. L. Brown-Triolo, M. J. Roach, K. Nelson, and R. J. Triolo, "Consumer perspectives on mobility: implications for neuroprosthesis design," J. Rehabil. Res. Dev., vol. 39, no. 6, pp. 659–669, Dec. 2002. + +[16] K. D. Anderson, "Consideration of user priorities when developing neural prosthetics," J. Neural Eng., vol. 6, no. 5, p. 055003, Sep. 2009, doi: 10.1088/1741-2560/6/5/055003. + +[17] A. G. Howard et al., "MobileNets: Efficient Convolutional Neural Networks for Mobile Vision Applications," ArXiv170404861 Cs, Apr. 2017, Accessed: May 21, 2021. [Online]. Available: http://arxiv.org/abs/1704.04861 + +[18] M. Sandler, A. Howard, M. Zhu, A. Zhmoginov, and L.-C. Chen, "MobileNetV2: Inverted Residuals and Linear Bottlenecks," ArXiv180104381 Cs, Mar. 2019, Accessed: May 21, 2021. [Online]. Available: http://arxiv.org/abs/1801.04381 + +[19] A. Howard et al., "Searching for MobileNetV3," ArXiv190502244 Cs, Nov. 2019, Accessed: May 21, 2021. [Online]. Available: http://arxiv.org/abs/1905.02244 + +[20] M. Tan and Q. V. Le, "EfficientNet: Rethinking Model Scaling for Convolutional Neural Networks," ArXiv190511946 Cs Stat, Sep. 2020, Accessed: May 21, 2021. [Online]. Available: http://arxiv.org/abs/1905.11946 + +[21] M. Tan and Q. V. Le, "EfficientNetV2: Smaller Models and Faster Training," ArXiv210400298 Cs, May 2021, Accessed: May 21, 2021. [Online]. Available: http://arxiv.org/abs/2104.00298 + +[22] H. Cai, C. Gan, T. Wang, Z. Zhang, and S. Han, "Once-for-All: Train One Network and Specialize it for Efficient Deployment," ArXiv190809791 Cs Stat, Apr. 2020, Accessed: May 21, 2021. [Online]. Available: http://arxiv.org/abs/1908.09791 + +[23] J. T. Springenberg, A. Dosovitskiy, T. Brox, and M. Riedmiller, "Striving for Simplicity: The All Convolutional Net," ArXiv14126806 Cs, Apr. 2015, Accessed: Jul. 23, 2021. [Online]. Available: http://arxiv.org/abs/1412.6806 + +[24] R. G. L. Koh, A. I. Nachman, and J. Zariffa, "Classification of naturally evoked compound action potentials in peripheral nerve spatiotemporal recordings," Sci. Rep., vol. 9, no. 1, p. 11145, Dec. 2019, doi: 10.1038/s41598-019-47450-8. + +[25] M. Lin, Q. Chen, and S. Yan, "Network In Network," ArXiv13124400 Cs, Mar. 2014, Accessed: Mar. 12, 2021. [Online]. Available: http://arxiv.org/abs/1312.4400 + +[26] K. Simonyan and A. Zisserman, "Very Deep Convolutional Networks for Large-Scale Image Recognition," ArXiv14091556 Cs, Apr. 2015, Accessed: Jun. 09, 2021. [Online]. Available: http://arxiv.org/abs/1409.1556 + +[27] N. Srivastava, G. Hinton, A. Krizhevsky, I. Sutskever, and R. Salakhutdinov, "Dropout: A Simple Way to Prevent Neural Networks from Overfitting," J. Mach. Learn. Res., vol. 15, no. 56, pp. 1929–1958, 2014. + +[28] C. Szegedy et al., "Going Deeper with Convolutions." arXiv, Sep. 16, 2014. Accessed: Jun. 05, 2022. [Online]. Available: http://arxiv.org/abs/1409.4842 + +[29] K. He, X. Zhang, S. Ren, and J. Sun, "Deep Residual Learning for Image Recognition," in 2016 IEEE Conference on Computer Vision and Pattern Recognition (CVPR), Las Vegas, NV, USA, Jun. 2016, pp. 770–778. doi: 10.1109/CVPR.2016.90. + +[30] Y.-H. Chen, T. Krishna, J. S. Emer, and V. Sze, "Eyeriss: An Energy-Efficient Reconfigurable Accelerator for Deep Convolutional Neural Networks," IEEE J. Solid-State Circuits, vol. 52, no. 1, pp. 127–138, Jan. 2017, doi: 10.1109/JSSC.2016.2616357. + +[31] J. Yue et al., "14.3 A 65nm Computing-in-Memory-Based CNN Processor with 2.9-to-35.8TOPS/W System Energy Efficiency Using Dynamic-Sparsity Performance-Scaling Architecture and Energy-Efficient Inter/Intra-Macro Data Reuse," in 2020 IEEE International Solid-State Circuits Conference - (ISSCC), Feb. 2020, pp. 234–236. doi: 10.1109/ISSCC19947.2020.9062958. + +[32] B. Moons and M. Verhelst, "A 0.3–2.6 TOPS/W precision-scalable processor for real-time large-scale ConvNets," in 2016 IEEE Symposium on VLSI Circuits (VLSI-Circuits), Jun. 2016, pp. 1–2. doi: 10.1109/VLSIC.2016.7573525. + +[33] X. Qiao et al., "A 65 nm 73 kb SRAM-Based Computing-In-Memory Macro With Dynamic-Sparsity Controlling," IEEE Trans. Circuits Syst. II Express Briefs, vol. 69, no. 6, pp. 2977–2981, Jun. 2022, doi: 10.1109/TCSII.2022.3162017. + +[34] I. Palit, Q. Lou, N. Acampora, J. Nahas, M. Niemier, and X. S. Hu, "Analytically modeling power and performance of a CNN system," in 2015 IEEE/ACM International Conference on Computer-Aided Design (ICCAD), Nov. 2015, pp. 186–193. doi: 10.1109/ICCAD.2015.7372569. diff --git a/scripts/Arduino/ArduinoOnDeviceTrain/ArduinoOnDeviceTrain.ino b/scripts/Arduino/ArduinoOnDeviceTrain/ArduinoOnDeviceTrain.ino new file mode 100644 index 0000000..4730071 --- /dev/null +++ b/scripts/Arduino/ArduinoOnDeviceTrain/ArduinoOnDeviceTrain.ino @@ -0,0 +1,275 @@ +/* + ArduinoOnDeviceTrain.ino + + Detects CAPs live from the streamed nerve signal and TRAINS a small model + on-device (rather than running a frozen one). Task: pricking vs movement + (movement = dorsiflexion + plantarflexion). + + For each detected CAP we extract 113 features, standardise them with running + stats, predict, and then do one SGD step using the label that came with the + stream. So the model learns while the data streams in. + + Stream protocol (same as stream_raw.py, 119 bytes): + [0xAA][0xBB] sync + [label] 1 byte: 0=dorsi 1=plantar 2=pricking 254=eval-only 255=unknown + [ts] 4 bytes uint32 LE + [ch0..ch55] 56 x uint16 LE raw Intan ADC (uV = 0.195*(raw-32768)) + + Per CAP it prints: + CAP ts=N pred=P label=Y p=PROB acc=ACC + pred/label: 1=pricking 0=movement (label=-1 if unknown) + acc: running accuracy over labelled CAPs so far + + LED: cyan idle, yellow armed, magenta computing, blue=pricking, green=movement. +*/ +#include +#include +#include +#include +#include "ondevice_model.h" + +// ── LED ────────────────────────────────────────────────────────────────────── +#define PHASE_IDLE 0 // cyan +#define PHASE_ARMED 1 // yellow +#define PHASE_COMPUTE 2 // magenta +#define PHASE_MOVEMENT 3 // green +#define PHASE_PRICKING 4 // blue + +static uint32_t result_hold_until = 0; + +static void set_led(int phase) { + bool r = false, g = false, b = false; + switch (phase) { + case PHASE_IDLE: g = true; b = true; break; // cyan + case PHASE_ARMED: r = true; g = true; break; // yellow + case PHASE_COMPUTE: r = true; b = true; break; // magenta + case PHASE_MOVEMENT: g = true; break; // green + case PHASE_PRICKING: b = true; break; // blue + } + digitalWrite(LEDR, r ? LOW : HIGH); + digitalWrite(LEDG, g ? LOW : HIGH); + digitalWrite(LEDB, b ? LOW : HIGH); +} + +// ── Signal processing (live CAP detection) ─────────────────────────────────── +#define N_CH 56 +#define N_RINGS 7 +#define CIRC_SIZE 150 // >= 100 + +// 2 kHz notch, Q=30, fs=30 kHz: [b0,b1,b2,a1,a2] +static const float NOTCH[5] = { + 0.9930669720f, -1.8144236429f, 0.9930669720f, + -1.8144236429f, 0.9861339440f +}; +// 1-3 kHz bandpass, 6th-order Butterworth, fs=30 kHz (6 SOS sections) +static const float BP[6][5] = { + { 0.0000407482f, 0.0000814964f, 0.0000407482f, -1.4614934920f, 0.6136710879f }, + { 1.0000000000f, 2.0000000000f, 1.0000000000f, -1.4091053901f, 0.6623962038f }, + { 1.0000000000f, 2.0000000000f, 1.0000000000f, -1.6315199139f, 0.7102771058f }, + { 1.0000000000f, -2.0000000000f, 1.0000000000f, -1.7840692078f, 0.8352552241f }, + { 1.0000000000f, -2.0000000000f, 1.0000000000f, -1.5099810120f, 0.8562240048f }, + { 1.0000000000f, -2.0000000000f, 1.0000000000f, -1.9018207709f, 0.9454426128f }, +}; + +static float circ_buf[N_CH][CIRC_SIZE]; +static int circ_write = 0; +static float notch_s[N_RINGS][2]; +static float bp_s[N_RINGS][6][2]; +static float daa_hist[8][N_RINGS]; +static int daa_write = 0; +static float ema_abs = 5.0f; +#define EMA_ALPHA 0.005f +static int refractory = 0; +static int countdown = -1; +static uint32_t peak_ts = 0; +static uint8_t peak_label = 255; + +// ── Model + running stats ──────────────────────────────────────────────────── +static RunningStandardizer od_std; +static LogRegSGD od_model; +static float od_window[OD_N_CH * OD_WIN]; // 56*100, static (22 KB) +static long seen = 0; // labelled CAPs +static long correct = 0; + +static inline float biquad(float x, const float* c, float* s) { + float w = x - c[3]*s[0] - c[4]*s[1]; + float y = c[0]*w + c[1]*s[0] + c[2]*s[1]; + s[1] = s[0]; + s[0] = w; + return y; +} + +static void classify_cap(uint32_t ts, uint8_t label); + +static void process_sample(const uint16_t* raw, uint32_t ts, uint8_t label) { + // 1. ADC -> uV + float uv[N_CH]; + for (int c = 0; c < N_CH; c++) uv[c] = 0.195f * ((float)raw[c] - 32768.f); + + // 2. Tripolar reference (mean of outer rings 0 and 6) + float ref = 0.f; + for (int c = 0; c < 8; c++) ref += uv[c]; + for (int c = 48; c < 56; c++) ref += uv[c]; + ref /= 16.f; + for (int c = 0; c < N_CH; c++) uv[c] -= ref; + + // 3. Circular buffer of referenced signal (for window extraction) + for (int c = 0; c < N_CH; c++) circ_buf[c][circ_write] = uv[c]; + circ_write = (circ_write + 1) % CIRC_SIZE; + + // 4. Ring averages + float ring_avg[N_RINGS]; + for (int r = 0; r < N_RINGS; r++) { + float s = 0.f; + for (int c = 0; c < 8; c++) s += uv[r*8 + c]; + ring_avg[r] = s * 0.125f; + } + + // 5. Causal notch + bandpass on each ring average + float filt[N_RINGS]; + for (int r = 0; r < N_RINGS; r++) { + float x = biquad(ring_avg[r], NOTCH, notch_s[r]); + for (int s = 0; s < 6; s++) x = biquad(x, BP[s], bp_s[r][s]); + filt[r] = x; + } + + // 6. Delay-and-add + for (int r = 0; r < N_RINGS; r++) daa_hist[daa_write][r] = filt[r]; + float daa = 0.f; + for (int r = 0; r < N_RINGS; r++) { + int hidx = (daa_write - r + 8) % 8; + daa += daa_hist[hidx][r]; + } + daa_write = (daa_write + 1) % 8; + + // 7. EMA noise floor -> adaptive threshold (~4 sigma) + float abs_daa = fabsf(daa); + ema_abs = (1.f - EMA_ALPHA) * ema_abs + EMA_ALPHA * abs_daa; + float threshold = 5.01f * ema_abs; + + // 8. Refractory + if (refractory > 0) refractory--; + + // 9. Threshold crossing -> arm + if (countdown < 0 && refractory == 0 && abs_daa > threshold) { + countdown = 50; + peak_ts = ts; + peak_label = label; + refractory = 150; + set_led(PHASE_ARMED); + } + + // 10. Countdown -> extract window and train/predict + if (countdown > 0) { + countdown--; + } else if (countdown == 0) { + countdown = -1; + classify_cap(peak_ts, peak_label); + } + + // 11. Back to idle when the result colour expires + if (countdown < 0 && (uint32_t)millis() >= result_hold_until) set_led(PHASE_IDLE); +} + +static void classify_cap(uint32_t ts, uint8_t label) { + set_led(PHASE_COMPUTE); + + // Build the 56x100 window (row-major c*WIN+t, t=0 oldest .. 99 newest) + for (int t = 0; t < OD_WIN; t++) { + int pos = (circ_write - OD_WIN + t + 2 * CIRC_SIZE) % CIRC_SIZE; + for (int c = 0; c < N_CH; c++) od_window[c * OD_WIN + t] = circ_buf[c][pos]; + } + + // label 254 = evaluate only: predict + report, but do NOT train or touch the + // running stats. Lets the host measure the frozen model on a held-out set so an + // adaptation drop/recovery is visible (the running accuracy alone hides it). + bool eval = (label == 254); + + // Features -> standardise -> predict + float feat[OD_N_FEAT]; + double z[OD_N_FEAT]; + od_extract_features(od_window, feat); + if (!eval) od_std.update(feat); + od_std.transform(feat, z); + double p = od_model.predict_proba(z); + int pred = (p >= 0.5) ? 1 : 0; + + set_led(pred == 1 ? PHASE_PRICKING : PHASE_MOVEMENT); + result_hold_until = (uint32_t)millis() + 300; + + if (eval) { + Serial.print("EVAL ts="); Serial.print(ts); + Serial.print(" pred="); Serial.print(pred); + Serial.print(" p="); Serial.println(p, 3); + Serial.flush(); + return; + } + + // Label: pricking(2)=1, dorsi/plantar(0,1)=0, unknown(255)=-1 + int y = (label == 255) ? -1 : (label == 2 ? 1 : 0); + + // Train (only if we have a label) and track running accuracy + if (y >= 0) { + seen++; + if (pred == y) correct++; + od_model.update(z, y); + } + + Serial.print("CAP ts="); Serial.print(ts); + Serial.print(" pred="); Serial.print(pred); + Serial.print(" label="); Serial.print(y); + Serial.print(" p="); Serial.print(p, 3); + Serial.print(" acc="); Serial.println(seen ? (double)correct / (double)seen : 0.0, 3); + Serial.flush(); +} + +// ── Serial protocol (same as stream_raw.py) ────────────────────────────────── +#define BAUD 115200 +#define PAYLOAD 117 // label(1) + ts(4) + 56*uint16(112) + +static uint8_t pkt_buf[PAYLOAD]; +static int pkt_pos = 0; +static int pkt_state = 0; // 0=wait 0xAA 1=wait 0xBB 2=payload + +void setup() { + pinMode(LEDR, OUTPUT); + pinMode(LEDG, OUTPUT); + pinMode(LEDB, OUTPUT); + set_led(PHASE_IDLE); + + Serial.begin(BAUD); + while (!Serial) {} + + memset(circ_buf, 0, sizeof(circ_buf)); + memset(notch_s, 0, sizeof(notch_s)); + memset(bp_s, 0, sizeof(bp_s)); + memset(daa_hist, 0, sizeof(daa_hist)); + + od_std.init(); + od_model.init(); // lr=0.05, l2=1e-4 + + Serial.println("ArduinoOnDeviceTrain ready."); +} + +void loop() { + while (Serial.available()) { + uint8_t b = (uint8_t)Serial.read(); + if (pkt_state == 0) { + if (b == 0xAA) pkt_state = 1; + } else if (pkt_state == 1) { + if (b == 0xBB) { pkt_state = 2; pkt_pos = 0; } + else if (b == 0xAA) pkt_state = 1; + else pkt_state = 0; + } else { + pkt_buf[pkt_pos++] = b; + if (pkt_pos == PAYLOAD) { + uint8_t label = pkt_buf[0]; + uint32_t ts; memcpy(&ts, pkt_buf + 1, 4); + uint16_t raw[N_CH]; memcpy(raw, pkt_buf + 5, N_CH * 2); + process_sample(raw, ts, label); + pkt_state = 0; + pkt_pos = 0; + } + } + } +} diff --git a/scripts/Arduino/ArduinoOnDeviceTrain/README.md b/scripts/Arduino/ArduinoOnDeviceTrain/README.md new file mode 100644 index 0000000..102d778 --- /dev/null +++ b/scripts/Arduino/ArduinoOnDeviceTrain/README.md @@ -0,0 +1,34 @@ +# ArduinoOnDeviceTrain + +Detects CAPs live from the streamed nerve data and **trains** a small model on +the board (pricking vs movement). It runs the full detection front-end and then +does one SGD step per detected CAP, using the label that comes with the stream. + +## Files +- `ArduinoOnDeviceTrain.ino` — the sketch (detection + on-device training) +- `ondevice_model.h` / `ondevice_model.cpp` — the model (features, standardizer, + logistic regression + SGD). Same code the host test in `../ondevice_train/` checks. + +## Run it +1. Open the sketch in the Arduino IDE and upload to the Nano 33 BLE Sense. +2. Stream the data to it from the laptop: + ``` + python scripts/cap_pipeline/ondevice/stream_raw.py + ``` +3. Watch the serial monitor. Per detected CAP it prints: + ``` + CAP ts=.. pred=.. label=.. p=.. acc=.. + ``` + `pred`/`label`: 1=pricking, 0=movement. `acc` is the running accuracy, which + should climb as it learns. + +LED: cyan = idle, yellow = CAP detected, magenta = computing, blue = pricking, +green = movement. + +## Notes +- Stream slower than real time if the board can't keep up (30 kHz x 56 ch is a lot + over serial) — it's a replay, so the rate doesn't have to be realtime. +- `stream_raw.py` drives the adaptation demo: it streams a sequence of `PHASES` + (e.g. RAT4 then RAT9), parses the `pred=`/`label=` reply, prints a sliding-window + accuracy so the shift and recovery are visible, and logs the learning curve to + `caches/learning_curve.csv`. Edit `PHASES`/`SERIAL_PORT` at the top of the file. diff --git a/scripts/Arduino/ArduinoOnDeviceTrain/ondevice_model.cpp b/scripts/Arduino/ArduinoOnDeviceTrain/ondevice_model.cpp new file mode 100644 index 0000000..b463ca1 --- /dev/null +++ b/scripts/Arduino/ArduinoOnDeviceTrain/ondevice_model.cpp @@ -0,0 +1,82 @@ +/* + ondevice_model.cpp -- see ondevice_model.h +*/ +#include "ondevice_model.h" +#include + +void od_extract_features(const float* window, float* out) { + float g_sum2 = 0.0f; + float g_min = window[0]; + float g_max = window[0]; + for (int c = 0; c < OD_N_CH; c++) { + const float* row = window + c * OD_WIN; + float sum2 = 0.0f, mn = row[0], mx = row[0]; + for (int t = 0; t < OD_WIN; t++) { + float v = row[t]; + sum2 += v * v; + if (v < mn) mn = v; + if (v > mx) mx = v; + g_sum2 += v * v; + if (v < g_min) g_min = v; + if (v > g_max) g_max = v; + } + out[c] = sqrtf(sum2 / OD_WIN); // per-channel RMS energy + out[OD_N_CH + c] = mx - mn; // per-channel peak-to-peak + } + out[2 * OD_N_CH] = sqrtf(g_sum2 / (OD_N_CH * OD_WIN)); // global energy + out[2 * OD_N_CH + 1] = g_max - g_min; // global p2p +} + +// ── RunningStandardizer (Welford) ─────────────────────────────────────────── +void RunningStandardizer::init() { + n = 0; + for (int i = 0; i < OD_N_FEAT; i++) { mean[i] = 0.0; M2[i] = 1.0; } +} + +void RunningStandardizer::update(const float* x) { + n++; + for (int i = 0; i < OD_N_FEAT; i++) { + double d = (double)x[i] - mean[i]; + mean[i] += d / (double)n; + M2[i] += d * ((double)x[i] - mean[i]); + } +} + +void RunningStandardizer::transform(const float* x, double* out) { + long denom = (n < 1) ? 1 : n; + for (int i = 0; i < OD_N_FEAT; i++) { + double sd = sqrt(M2[i] / (double)denom) + 1e-6; + out[i] = ((double)x[i] - mean[i]) / sd; + } +} + +// ── LogRegSGD ─────────────────────────────────────────────────────────────── +void LogRegSGD::init(double lr_, double l2_) { + for (int i = 0; i < OD_N_FEAT; i++) w[i] = 0.0; + b = 0.0; lr = lr_; l2 = l2_; +} + +static inline double od_sigmoid(double z) { + if (z > 30.0) z = 30.0; + if (z < -30.0) z = -30.0; + return 1.0 / (1.0 + exp(-z)); +} + +double LogRegSGD::predict_proba(const double* x) const { + double z = b; + for (int i = 0; i < OD_N_FEAT; i++) z += w[i] * x[i]; + return od_sigmoid(z); +} + +int LogRegSGD::predict(const double* x) const { + return predict_proba(x) >= 0.5 ? 1 : 0; +} + +double LogRegSGD::update(const double* x, int y) { + double p = predict_proba(x); + double g = p - (double)y; + for (int i = 0; i < OD_N_FEAT; i++) w[i] -= lr * (g * x[i] + l2 * w[i]); + b -= lr * g; + const double eps = 1e-9; + return -((double)y * log(p + eps) + (1.0 - (double)y) * log(1.0 - p + eps)); +} diff --git a/scripts/Arduino/ArduinoOnDeviceTrain/ondevice_model.h b/scripts/Arduino/ArduinoOnDeviceTrain/ondevice_model.h new file mode 100644 index 0000000..155e2db --- /dev/null +++ b/scripts/Arduino/ArduinoOnDeviceTrain/ondevice_model.h @@ -0,0 +1,43 @@ +/* + ondevice_model.h + + The model we train on-device: feature extraction, a running standardizer and + logistic regression with SGD. Plain C++/math.h (no Arduino.h) so it builds on + the host for testing and on the Nano in a sketch. + + Mirrors scripts/cap_pipeline/ondevice/model.py. + Task: pricking (1) vs movement (0). +*/ +#ifndef ONDEVICE_MODEL_H +#define ONDEVICE_MODEL_H + +#define OD_N_CH 56 +#define OD_WIN 100 +#define OD_N_FEAT (OD_N_CH * 2 + 2) // 113 + +// window: 56*100 floats, row-major (channel*WIN + t). out: 113 floats. +void od_extract_features(const float* window, float* out); + +struct RunningStandardizer { + long n; + double mean[OD_N_FEAT]; + double M2[OD_N_FEAT]; + + void init(); + void update(const float* x); // feed one feature vector + void transform(const float* x, double* out); // z-score with current stats +}; + +struct LogRegSGD { + double w[OD_N_FEAT]; + double b; + double lr; + double l2; + + void init(double lr = 0.05, double l2 = 1e-4); + double predict_proba(const double* x) const; + int predict(const double* x) const; + double update(const double* x, int y); // one SGD step, returns BCE loss +}; + +#endif diff --git a/scripts/Arduino/ondevice_train/.gitignore b/scripts/Arduino/ondevice_train/.gitignore new file mode 100644 index 0000000..f610049 --- /dev/null +++ b/scripts/Arduino/ondevice_train/.gitignore @@ -0,0 +1,4 @@ +# generated by make_parity.py / the build +*.bin +meta.txt +test_host diff --git a/scripts/Arduino/ondevice_train/README.md b/scripts/Arduino/ondevice_train/README.md new file mode 100644 index 0000000..c4b97eb --- /dev/null +++ b/scripts/Arduino/ondevice_train/README.md @@ -0,0 +1,16 @@ +# Host test for the on-device model + +Checks that the C++ model used by the sketch matches the Python reference +(`scripts/cap_pipeline/ondevice/model.py`). The model itself lives in the sketch +folder `../ArduinoOnDeviceTrain/` (`ondevice_model.h` / `.cpp`) so there's one copy. + +- `make_parity.py` — exports test data + the Python results to compare against. +- `test_host.cpp` — runs the C++ model on that data and checks it matches Python. + +## Run +``` +python make_parity.py +g++ -O2 -std=c++17 test_host.cpp ../ArduinoOnDeviceTrain/ondevice_model.cpp -o test_host +./test_host +``` +Should print all PASS (predictions identical, weights match to ~1e-15). diff --git a/scripts/Arduino/ondevice_train/make_parity.py b/scripts/Arduino/ondevice_train/make_parity.py new file mode 100644 index 0000000..e3def27 --- /dev/null +++ b/scripts/Arduino/ondevice_train/make_parity.py @@ -0,0 +1,78 @@ +""" +make_parity.py + +Export test data + expected results so the C++ training core (ondevice_model.*) +can be checked against the Python reference (model.py). Run this, then build and +run test_host.cpp. + +Writes (all little-endian raw dumps, read back in test_host.cpp): + meta.txt "N K" (N streaming samples, K windows for the feature test) + raw_windows.bin float32 K x 56 x 100 windows for the feature-parity check + py_features.bin float32 K x 113 Python extract_features(window) + features.bin float32 N x 113 features for the streaming check + labels.bin uint8 N 0=movement, 1=pricking + py_preds.bin uint8 N Python prequential predictions + expected.bin double [acc, b, w0..w112] Python final accuracy + weights +""" +import sys, pathlib +import numpy as np + +HERE = pathlib.Path(__file__).resolve().parent +OND = HERE.parents[1] / "cap_pipeline" / "ondevice" +CACHE = HERE.parents[1] / "cap_pipeline" / "caches" / "cap_windows_rich.npz" +sys.path.insert(0, str(OND)) +from model import extract_features, RunningStandardizer, LogRegSGD, N_FEAT + +K = 50 # windows for the feature-parity check + + +def main(): + if not CACHE.exists(): + raise SystemExit(f"Missing {CACHE}. Build it: python scripts/cap_pipeline/ondevice/build_dataset.py") + d = np.load(CACHE) + X, y3, rat = d["X"], d["y"], d["rat"] + y = (y3 == 2).astype(np.int64) # 1 = pricking, 0 = movement + + # Deterministic balanced RAT4 subset (matches showcase's balanced() with seed 0) + rng = np.random.default_rng(0) + idx = np.where(rat == 4)[0] + i0, i1 = idx[y[idx] == 0], idx[y[idx] == 1] + n = min(len(i0), len(i1)) + order = rng.permutation(np.concatenate([rng.choice(i0, n, False), rng.choice(i1, n, False)])) + Xs, ys = X[order], y[order] + N = len(Xs) + + # Features for every streaming sample (float32, exactly as model.py produces) + feats = np.stack([extract_features(Xs[i]) for i in range(N)]).astype(np.float32) + + # Reference streaming pass (same loop as showcase, default LogRegSGD) + model, std = LogRegSGD(N_FEAT), RunningStandardizer(N_FEAT) + preds = np.empty(N, dtype=np.uint8) + correct = 0 + for i in range(N): + f = feats[i] + std.update(f) + fz = std.transform(f) + p = model.predict(fz) + preds[i] = p + correct += (p == ys[i]) + model.update(fz, ys[i]) + acc = correct / N + + # Write everything + (HERE / "meta.txt").write_text(f"{N} {K}\n") + Xs[:K].astype(np.float32).tofile(HERE / "raw_windows.bin") + feats[:K].astype(np.float32).tofile(HERE / "py_features.bin") + feats.astype(np.float32).tofile(HERE / "features.bin") + ys.astype(np.uint8).tofile(HERE / "labels.bin") + preds.tofile(HERE / "py_preds.bin") + expected = np.concatenate([[acc, model.b], model.w]).astype(np.float64) + expected.tofile(HERE / "expected.bin") + + print(f"N={N} streaming samples, K={K} feature-test windows") + print(f"Python prequential accuracy = {acc:.4f}") + print(f"Wrote parity files to {HERE}") + + +if __name__ == "__main__": + main() diff --git a/scripts/Arduino/ondevice_train/test_host.cpp b/scripts/Arduino/ondevice_train/test_host.cpp new file mode 100644 index 0000000..78fcaa3 --- /dev/null +++ b/scripts/Arduino/ondevice_train/test_host.cpp @@ -0,0 +1,108 @@ +/* + test_host.cpp + + Checks the C++ training core (ondevice_model.*) against the Python reference. + Run make_parity.py first to produce the data files, then: + + g++ -O2 -std=c++17 test_host.cpp ondevice_model.cpp -o test_host && ./test_host + + Reads the .bin files in this folder, reruns the same feature extraction and + streaming SGD in C++, and compares to the exported Python results. +*/ +#include "../ArduinoOnDeviceTrain/ondevice_model.h" +#include +#include +#include +#include +#include +#include + +static std::vector read_f32(const std::string& p) { + FILE* f = fopen(p.c_str(), "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", p.c_str()); exit(1); } + fseek(f, 0, SEEK_END); long n = ftell(f) / 4; fseek(f, 0, SEEK_SET); + std::vector v(n); size_t r = fread(v.data(), 4, n, f); (void)r; fclose(f); + return v; +} +static std::vector read_u8(const std::string& p) { + FILE* f = fopen(p.c_str(), "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", p.c_str()); exit(1); } + fseek(f, 0, SEEK_END); long n = ftell(f); fseek(f, 0, SEEK_SET); + std::vector v(n); size_t r = fread(v.data(), 1, n, f); (void)r; fclose(f); + return v; +} +static std::vector read_f64(const std::string& p) { + FILE* f = fopen(p.c_str(), "rb"); + if (!f) { fprintf(stderr, "cannot open %s\n", p.c_str()); exit(1); } + fseek(f, 0, SEEK_END); long n = ftell(f) / 8; fseek(f, 0, SEEK_SET); + std::vector v(n); size_t r = fread(v.data(), 8, n, f); (void)r; fclose(f); + return v; +} + +int main() { + long N = 0, K = 0; + { FILE* f = fopen("meta.txt", "r"); if (!f) { fprintf(stderr, "run make_parity.py first\n"); return 1; } + if (fscanf(f, "%ld %ld", &N, &K) != 2) { fprintf(stderr, "bad meta.txt\n"); return 1; } fclose(f); } + + bool ok = true; + + // ── Test A: feature extraction parity ─────────────────────────────────── + std::vector raw = read_f32("raw_windows.bin"); // K x 56 x 100 + std::vector pyf = read_f32("py_features.bin"); // K x 113 + double feat_maxdiff = 0.0; + std::vector cf(OD_N_FEAT); + for (long k = 0; k < K; k++) { + od_extract_features(raw.data() + k * OD_N_CH * OD_WIN, cf.data()); + for (int i = 0; i < OD_N_FEAT; i++) { + double dd = fabs((double)cf[i] - (double)pyf[k * OD_N_FEAT + i]); + if (dd > feat_maxdiff) feat_maxdiff = dd; + } + } + bool feat_ok = feat_maxdiff < 1e-3; + ok &= feat_ok; + printf("Test A features : max|C++ - Py| = %.3e [%s]\n", + feat_maxdiff, feat_ok ? "PASS" : "FAIL"); + + // ── Test B: streaming SGD parity ──────────────────────────────────────── + std::vector feats = read_f32("features.bin"); // N x 113 + std::vector labels = read_u8("labels.bin"); // N + std::vector pypred = read_u8("py_preds.bin"); // N + std::vector exp = read_f64("expected.bin"); // [acc, b, w0..w112] + + RunningStandardizer std_; std_.init(); + LogRegSGD model; model.init(); + std::vector z(OD_N_FEAT); + long correct = 0, pred_mismatch = 0; + for (long i = 0; i < N; i++) { + const float* f = feats.data() + i * OD_N_FEAT; + std_.update(f); + std_.transform(f, z.data()); + int p = model.predict(z.data()); + if (p != (int)pypred[i]) pred_mismatch++; + if (p == (int)labels[i]) correct++; + model.update(z.data(), (int)labels[i]); + } + double cpp_acc = (double)correct / (double)N; + double py_acc = exp[0], py_b = exp[1]; + double acc_diff = fabs(cpp_acc - py_acc); + double w_maxdiff = fabs(model.b - py_b); + for (int i = 0; i < OD_N_FEAT; i++) { + double dd = fabs(model.w[i] - exp[2 + i]); + if (dd > w_maxdiff) w_maxdiff = dd; + } + bool pred_ok = (pred_mismatch == 0); + bool acc_ok = (acc_diff < 1e-9); + bool w_ok = (w_maxdiff < 1e-6); + ok &= pred_ok & acc_ok & w_ok; + + printf("Test B predictions : %ld / %ld differ from Python [%s]\n", + pred_mismatch, N, pred_ok ? "PASS" : "FAIL"); + printf("Test B accuracy : C++ %.4f Py %.4f (diff %.2e) [%s]\n", + cpp_acc, py_acc, acc_diff, acc_ok ? "PASS" : "FAIL"); + printf("Test B weights : max|C++ - Py| = %.3e [%s]\n", + w_maxdiff, w_ok ? "PASS" : "FAIL"); + + printf("\n%s\n", ok ? "ALL PASS - C++ core matches the Python reference" + : "FAIL - see above"); + return ok ? 0 : 1; +} diff --git a/scripts/cap_pipeline/ondevice/README.md b/scripts/cap_pipeline/ondevice/README.md new file mode 100644 index 0000000..a196a6f --- /dev/null +++ b/scripts/cap_pipeline/ondevice/README.md @@ -0,0 +1,36 @@ +# On-device part: pricking vs movement + +This is the on-device training part. We train a small model on the nerve-cuff +data and classify pricking vs movement (movement = dorsiflexion + plantarflexion). +It's written in plain numpy here so we can port it to Arduino C++ later. + +We use this task because dorsi vs plantar basically can't be told apart in this +data (~50%, see `docs/replication_investigation_2026-06.md`), but pricking vs +movement works ok — around 70% when we stay within one rat. + +## Files +- `build_dataset.py` — reads the raw `.rhd` files, detects CAPs, saves the windows +- `model.py` — the model: feature extraction, a running standardizer, logistic regression + SGD +- `showcase.py` — runs the demos (online learning + adaptation) + +## How to run +Build the window cache once (takes a couple of minutes): + + python build_dataset.py + +Then run the demos: + + python showcase.py + +The raw data goes in `/data` (use `data/download_dataset.py` to get it). + +## What the demos do +- **Online learning** — feed windows one at a time, accuracy goes up as it learns. +- **Cross-rat** — train on a few rats, accuracy drops on a new rat, then we keep + training on the new rat to see if it comes back up. +- **Channel dropout** — zero out some electrodes (like they broke), accuracy drops, + then the model adapts. + +## Next +Port `model.py` to Arduino C++ (features, standardizer, logistic regression + SGD) +and feed it the replay stream. diff --git a/scripts/cap_pipeline/ondevice/build_dataset.py b/scripts/cap_pipeline/ondevice/build_dataset.py new file mode 100644 index 0000000..b85f0ec --- /dev/null +++ b/scripts/cap_pipeline/ondevice/build_dataset.py @@ -0,0 +1,204 @@ +""" +build_dataset.py + +Reads the raw .rhd recordings, detects CAPs and saves them as windows we can +train on. Run this once before showcase.py. + +Detection per file: tripolar reference -> notch (2kHz) + bandpass (1-3kHz) on the +ring averages -> delay-and-add -> threshold on that -> drop big artefacts (>15uV) +-> cut a 56x100 window around each detected CAP (49 before, peak, 50 after). + +Output: ../caches/cap_windows_rich.npz + X : (N, 56, 100) float32 windows in uV + y : (N,) 0=dorsi, 1=plantar, 2=pricking + rat : (N,) rat number + +Usage: + python build_dataset.py build if the cache isn't there + python build_dataset.py --force rebuild +""" +import argparse, pathlib, struct, time +import numpy as np +from scipy.signal import sosfilt, lfilter, lfilter_zi, butter, iirnotch, tf2sos + +HERE = pathlib.Path(__file__).resolve().parent +DATA_DIR = HERE.parents[2] / "data" # /data +CACHE = HERE.parent / "caches" / "cap_windows_rich.npz" + +# Each recording folder and its class label. (RAT7 uses a 64-ch cuff -> skipped.) +STREAM_FILES = [ + ("RAT3_PRICKING", 2), + ("RAT4_DORSIFLEXION", 0), ("RAT4_PLANTARFLEXION", 1), ("RAT4_PRICKING", 2), + ("RAT5_DORSIFLEXION", 0), ("RAT5_PLANTARFLEXION", 1), ("RAT5_PRICKING", 2), + ("RAT6_DORSIFLEXION", 0), ("RAT6_PLANTARFLEXION", 1), ("RAT6_PRICKING", 2), + ("RAT7_DORSIFLEXION", 0), ("RAT7_PLANTARFLEXION", 1), ("RAT7_PRICKING", 2), + ("RAT8_DORSIFLEXION", 0), ("RAT8_PLANTARFLEXION", 1), ("RAT8_PRICKING", 2), + ("RAT9_DORSIFLEXION", 0), ("RAT9_PLANTARFLEXION", 1), ("RAT9_PRICKING", 2), + ("RAT10_DORSIFLEXION", 0), ("RAT10_PLANTARFLEXION", 1), ("RAT10_PRICKING", 2), +] +CLASS_NAMES = ["Dorsiflexion", "Plantarflexion", "Pricking"] + +# Detection parameters +FS, N_CH, N_RINGS = 30_000.0, 56, 7 +EMA_ALPHA, EMA_INIT = 0.005, 5.0 +THRESHOLD, REFRACTORY = 5.01, 150 +WIN_SAMPS = 100 +WARMUP = int(0.5 * FS) +PEAK_UV = 15.0 # discard windows whose DAA peak exceeds this (artifact reject) + +b_n, a_n = iirnotch(2000.0, 30.0, fs=FS) +SOS_NOTCH = tf2sos(b_n, a_n) +SOS_BP = butter(6, [1000.0, 3000.0], btype="band", fs=FS, output="sos") + + +# ── Intan .rhd reader ──────────────────────────────────────────────────────── +def _rqs(f): + r = f.read(4) + if len(r) < 4: return "" + n = struct.unpack("= 1) or major > 1: n_temp = struct.unpack("= 3) or major > 1: f.read(2) + if major > 1: _rqs(f) + ng = struct.unpack(" 0 and en: + for _ in range(nch): + _rqs(f); _rqs(f) + _, _, st, ce, _, _ = struct.unpack(" keep first 56 + return None + spb, n_amp = hdr["spb"], hdr["n_amp"] + b_ts, b_amp = spb * 4, spb * n_amp * 2 + b_skip = ((spb // 4) * hdr["n_aux"] * 2 + hdr["n_supply"] * 2 + + hdr["n_temp"] * 2 + spb * hdr["n_adc"] * 2 + + (spb * 2 if hdr["n_dig_in"] > 0 else 0) + + (spb * 2 if hdr["n_dig_out"] > 0 else 0)) + blocks = [] + while True: + if len(f.read(b_ts)) < b_ts: break + amp = f.read(b_amp) + if len(amp) < b_amp: break + blocks.append(np.frombuffer(amp, dtype=np.uint16).reshape(n_amp, spb, order="F")) + if b_skip and len(f.read(b_skip)) < b_skip: break + except Exception: + return None + if not blocks: + return None + raw = np.concatenate(blocks, axis=1)[:N_CH] # drop any extra channels (RAT7: 64 -> 56) + return (0.195 * (raw.astype(np.float32) - 32768.0)).T # (N_samples, N_CH) + + +# ── Vectorized CAP detection ───────────────────────────────────────────────── +def collect_windows(rhd_paths): + """Detect CAPs and return a list of (56, 100) float32 windows.""" + windows = [] + for path in rhd_paths: + uv = _read_rhd(path) + if uv is None: + continue + N = len(uv) + ref = (uv[:, :8].mean(axis=1) + uv[:, 48:].mean(axis=1)) / 2.0 + uv -= ref[:, None] + rings = uv[:, :N_RINGS * 8].reshape(N, N_RINGS, 8).mean(axis=2).astype(np.float64) + for r in range(N_RINGS): + rings[:, r] = sosfilt(SOS_NOTCH, rings[:, r]) + rings[:, r] = sosfilt(SOS_BP, rings[:, r]) + daa = np.zeros(N, dtype=np.float64) + for r in range(N_RINGS): + if r == 0: daa += rings[:, r] + else: daa[r:] += rings[:-r, r] + b_ema = np.array([EMA_ALPHA]); a_ema = np.array([1.0, -(1.0 - EMA_ALPHA)]) + zi = lfilter_zi(b_ema, a_ema) * EMA_INIT + ema_arr, _ = lfilter(b_ema, a_ema, np.abs(daa), zi=zi) + abs_daa = np.abs(daa) + candidates = np.where(abs_daa[WARMUP:] > THRESHOLD * ema_arr[WARMUP:])[0] + WARMUP + next_ok = WARMUP + for trig in candidates: + if trig < next_ok: + continue + next_ok = trig + REFRACTORY + if float(abs_daa[trig]) > PEAK_UV: + continue + start, end = trig - 49, trig + 51 + if start < WARMUP or end > N: + continue + windows.append(uv[start:end, :].T.astype(np.float32).copy()) + return windows + + +def build(): + Xs, ys, rats = [], [], [] + for subdir, label in STREAM_FILES: + d = DATA_DIR / subdir + if not d.exists(): + print(f" skip {subdir} (folder missing)", flush=True); continue + paths = sorted(d.glob("*.rhd")) + if not paths: + print(f" skip {subdir} (no .rhd)", flush=True); continue + rat = int(subdir.split("_")[0].replace("RAT", "")) + wins = collect_windows(paths) + print(f" RAT{rat:<2d} {CLASS_NAMES[label]:14s}: {len(wins)} windows", flush=True) + if wins: + Xs.append(np.stack(wins)) + ys.append(np.full(len(wins), label, dtype=np.int64)) + rats.append(np.full(len(wins), rat, dtype=np.int64)) + if not Xs: + raise SystemExit(f"No windows detected. Is the dataset under {DATA_DIR}? " + f"(see data/download_dataset.py)") + return np.concatenate(Xs), np.concatenate(ys), np.concatenate(rats) + + +def main(): + ap = argparse.ArgumentParser(description=__doc__) + ap.add_argument("--force", action="store_true", help="rebuild even if cache exists") + args = ap.parse_args() + + if CACHE.exists() and not args.force: + d = np.load(CACHE) + print(f"Cache already exists: {CACHE} ({len(d['X'])} windows). Use --force to rebuild.") + return + + t0 = time.time() + print(f"Detecting CAP windows from {DATA_DIR} ...", flush=True) + X, y, rat = build() + CACHE.parent.mkdir(parents=True, exist_ok=True) + np.savez(CACHE, X=X, y=y, rat=rat) + print(f"\nSaved {len(X)} windows -> {CACHE}") + for c, name in enumerate(CLASS_NAMES): + print(f" {name:14s}: {(y == c).sum()}") + print(f"Done in {time.time() - t0:.0f}s") + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/model.py b/scripts/cap_pipeline/ondevice/model.py new file mode 100644 index 0000000..a495f0f --- /dev/null +++ b/scripts/cap_pipeline/ondevice/model.py @@ -0,0 +1,76 @@ +""" +model.py + +The model we want to run on the Arduino, kept in plain numpy so it's easy to +rewrite in C++ later. + +- extract_features: turn a 56x100 window into 113 features +- RunningStandardizer: running mean/std so we can scale features while streaming +- LogRegSGD: logistic regression, trained one sample at a time + +Task is pricking vs movement. We don't normalise each window on purpose - for +pricking the amplitude is the main thing that separates the classes. +""" +import numpy as np + +N_CH = 56 +WIN = 100 +N_FEAT = N_CH * 2 + 2 # per-channel energy + per-channel p2p + 2 globals = 113 + + +# ── Feature extraction ────────────────────────────────────────────────────── +def extract_features(window): + """(56, 100) window -> 113 features.""" + energy = np.sqrt((window ** 2).mean(axis=1)) # 56 per-channel RMS energy + p2p = window.max(axis=1) - window.min(axis=1) # 56 per-channel peak-to-peak + g_energy = np.sqrt((window ** 2).mean()) # 1 global RMS energy + g_p2p = window.max() - window.min() # 1 global peak-to-peak + return np.concatenate([energy, p2p, [g_energy, g_p2p]]).astype(np.float32) + + +# ── Online feature standardizer (Welford running mean/variance) ───────────── +class RunningStandardizer: + """Running z-score, updated as samples come in.""" + def __init__(self, n): + self.n = 0 + self.mean = np.zeros(n, np.float64) + self.M2 = np.ones(n, np.float64) # start at 1 -> unit variance before data + + def update(self, x): + self.n += 1 + d = x - self.mean + self.mean += d / self.n + self.M2 += d * (x - self.mean) + + def transform(self, x): + std = np.sqrt(self.M2 / max(1, self.n)) + 1e-6 + return (x - self.mean) / std + + +# ── Binary logistic regression trained with one-sample SGD ────────────────── +class LogRegSGD: + """Logistic regression (weights + bias), trained one sample at a time.""" + def __init__(self, n_feat, lr=0.05, l2=1e-4): + self.w = np.zeros(n_feat, np.float64) + self.b = 0.0 + self.lr = lr + self.l2 = l2 + + @staticmethod + def _sigmoid(z): + return 1.0 / (1.0 + np.exp(-np.clip(z, -30, 30))) + + def predict_proba(self, x): + return self._sigmoid(self.w @ x + self.b) + + def predict(self, x): + return int(self.predict_proba(x) >= 0.5) + + def update(self, x, y): + """One SGD step on a single (x, y). Returns the per-sample BCE loss.""" + p = self.predict_proba(x) + g = p - y # dL/dz for sigmoid + BCE + self.w -= self.lr * (g * x + self.l2 * self.w) + self.b -= self.lr * g + eps = 1e-9 + return -(y * np.log(p + eps) + (1 - y) * np.log(1 - p + eps)) diff --git a/scripts/cap_pipeline/ondevice/model_cnn.py b/scripts/cap_pipeline/ondevice/model_cnn.py new file mode 100644 index 0000000..ec5ef0a --- /dev/null +++ b/scripts/cap_pipeline/ondevice/model_cnn.py @@ -0,0 +1,89 @@ +""" +model_cnn.py + +ESCAPE-NET-Mini-style two-stream CNN (Hwang et al. 2022, sec 2.4.1) on the raw +56x100 window. Two inputs: spatial-emphasis (rows grouped by ring) and +temporal-emphasis (rows grouped by line along the nerve). The window is globally +standardized, NOT per-window — amplitude is a real cue for pricking. + +Same task/split as the other models. Saves test probabilities for the ensemble. + + python model_cnn.py +""" +import numpy as np +import torch +import torch.nn as nn + +from task_data import load, splits, save_preds, report, SEED + +NAME = "cnn" +EPOCHS = 60 + + +def temporal_order(): + # ring-major (ring*8+contact) -> line-major (contact*7+ring) + return [ring * 8 + contact for contact in range(8) for ring in range(7)] + + +def inputs(X, idx, mean, std): + W = (X[idx] - mean) / std + perm = temporal_order() + sp = torch.tensor(W[:, None, :, :], dtype=torch.float32) + te = torch.tensor(W[:, perm][:, None, :, :], dtype=torch.float32) + return sp, te + + +class TwoStream(nn.Module): + def __init__(self): + super().__init__() + + def stream(): + return nn.Sequential( + nn.Conv2d(1, 4, 8), nn.ReLU(), nn.MaxPool2d(2), + nn.Conv2d(4, 4, 4), nn.ReLU(), nn.MaxPool2d(2), + nn.Conv2d(4, 4, 2), nn.ReLU(), + nn.Conv2d(4, 4, 1), nn.ReLU(), + nn.Flatten()) + self.a, self.b = stream(), stream() + self.head = nn.Sequential(nn.LazyLinear(8), nn.ReLU(), nn.Linear(8, 2)) + + def forward(self, xa, xb): + return self.head(torch.cat([self.a(xa), self.b(xb)], 1)) + + +def fit_predict(Xtr_sp, Xtr_te, ytr, Xte_sp, Xte_te): + torch.manual_seed(SEED) + net = TwoStream() + net(Xtr_sp[:2], Xtr_te[:2]) # init LazyLinear + opt = torch.optim.Adam(net.parameters(), lr=1e-3, weight_decay=1e-4) + lossf = nn.CrossEntropyLoss() + yt = torch.tensor(ytr, dtype=torch.long) + n, bs = len(ytr), 64 + for _ in range(EPOCHS): + net.train() + for b in torch.randperm(n).split(bs): + opt.zero_grad() + lossf(net(Xtr_sp[b], Xtr_te[b]), yt[b]).backward() + opt.step() + net.eval() + with torch.no_grad(): + return torch.softmax(net(Xte_sp, Xte_te), 1)[:, 1].numpy() + + +def main(): + X, _, _, y = load() + rats, probs, labels = [], [], [] + for r, tr, te in splits(): + mean, std = X[tr].mean(0), X[tr].std(0) + 1e-6 + sp_tr, te_tr = inputs(X, tr, mean, std) + sp_te, te_te = inputs(X, te, mean, std) + p = fit_predict(sp_tr, te_tr, y[tr], sp_te, te_te) + rats += [r] * len(te); probs += list(p); labels += list(y[te]) + print(f" RAT{r} done", flush=True) + + save_preds(NAME, rats, probs, labels) + report(NAME, rats, probs, labels) + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/model_ensemble.py b/scripts/cap_pipeline/ondevice/model_ensemble.py new file mode 100644 index 0000000..ff8d1b7 --- /dev/null +++ b/scripts/cap_pipeline/ondevice/model_ensemble.py @@ -0,0 +1,49 @@ +""" +model_ensemble.py + +Soft-voting ensemble: averages P(pricking) from the linear, random forest and CNN +models, then predicts the higher-probability class. Run the three model scripts +first (they save their test probabilities); this loads and combines them, and +prints all four side by side. + + python model_linear.py + python model_random_forest.py + python model_cnn.py + python model_ensemble.py +""" +import numpy as np + +from task_data import load_preds, report, RATS + +BASE = ["linear", "random_forest", "cnn"] + + +def main(): + loaded = {} + for name in BASE: + try: + loaded[name] = load_preds(name) + except FileNotFoundError: + raise SystemExit(f"Missing predictions for '{name}'. Run model_{name}.py first.") + + # all models share the same test set/order; verify via labels + ref_rat, _, ref_lab = loaded[BASE[0]] + for name, (rat, _, lab) in loaded.items(): + if not (np.array_equal(rat, ref_rat) and np.array_equal(lab, ref_lab)): + raise SystemExit(f"'{name}' test set doesn't match — re-run all model scripts.") + + probs = np.mean([loaded[n][1] for n in BASE], axis=0) # soft vote + report("ensemble (linear+rf+cnn)", ref_rat, probs, ref_lab) + + # side-by-side mean accuracy + print("\nmean accuracy over rats:") + for name in BASE: + rat, p, lab = loaded[name] + accs = [((p[rat == r] >= 0.5).astype(int) == lab[rat == r]).mean() for r in RATS] + print(f" {name:14s} {np.mean(accs):.3f}") + ens = [((probs[ref_rat == r] >= 0.5).astype(int) == ref_lab[ref_rat == r]).mean() for r in RATS] + print(f" {'ensemble':14s} {np.mean(ens):.3f}") + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/model_linear.py b/scripts/cap_pipeline/ondevice/model_linear.py new file mode 100644 index 0000000..f3df8fd --- /dev/null +++ b/scripts/cap_pipeline/ondevice/model_linear.py @@ -0,0 +1,32 @@ +""" +model_linear.py + +Logistic regression on the 113 hand features (the model we run on the Arduino). +Within-rat, temporal split. Saves test probabilities for the ensemble. + + python model_linear.py +""" +import numpy as np +from sklearn.linear_model import LogisticRegression +from sklearn.preprocessing import StandardScaler + +from task_data import load, splits, feats, save_preds, report + +NAME = "linear" + + +def main(): + X, _, _, y = load() + rats, probs, labels = [], [], [] + for r, tr, te in splits(): + sc = StandardScaler().fit(feats(X, tr)) + clf = LogisticRegression(max_iter=2000).fit(sc.transform(feats(X, tr)), y[tr]) + p = clf.predict_proba(sc.transform(feats(X, te)))[:, 1] + rats += [r] * len(te); probs += list(p); labels += list(y[te]) + + save_preds(NAME, rats, probs, labels) + report(NAME, rats, probs, labels) + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/model_random_forest.py b/scripts/cap_pipeline/ondevice/model_random_forest.py new file mode 100644 index 0000000..e292f9c --- /dev/null +++ b/scripts/cap_pipeline/ondevice/model_random_forest.py @@ -0,0 +1,30 @@ +""" +model_random_forest.py + +Random forest on the 113 hand features. Same task/split as model_linear.py so +the numbers are comparable. Saves test probabilities for the ensemble. + + python model_random_forest.py +""" +from sklearn.ensemble import RandomForestClassifier + +from task_data import load, splits, feats, save_preds, report, SEED + +NAME = "random_forest" + + +def main(): + X, _, _, y = load() + rats, probs, labels = [], [], [] + for r, tr, te in splits(): + clf = RandomForestClassifier(n_estimators=200, random_state=SEED, n_jobs=-1) + clf.fit(feats(X, tr), y[tr]) + p = clf.predict_proba(feats(X, te))[:, 1] + rats += [r] * len(te); probs += list(p); labels += list(y[te]) + + save_preds(NAME, rats, probs, labels) + report(NAME, rats, probs, labels) + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/showcase.py b/scripts/cap_pipeline/ondevice/showcase.py new file mode 100644 index 0000000..f40ce38 --- /dev/null +++ b/scripts/cap_pipeline/ondevice/showcase.py @@ -0,0 +1,134 @@ +""" +showcase.py + +Runs the model from model.py on the data to show the on-device learning idea. +Three demos: + + 1. online learning: train one sample at a time on one rat, watch accuracy climb + 2. cross-rat adaptation: train on some rats, test on a new one, then keep + training on the new rat and see if it recovers + 3. channel dropout: zero some electrodes, then let the model adapt + +Task: pricking vs movement. Needs the cache from build_dataset.py. + +Usage: + python showcase.py +""" +import sys, pathlib, time +import numpy as np + +HERE = pathlib.Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +from model import extract_features, RunningStandardizer, LogRegSGD, N_CH, N_FEAT + +CACHE = HERE.parent / "caches" / "cap_windows_rich.npz" + + +# ── Data helpers ───────────────────────────────────────────────────────────── +def load_task(): + if not CACHE.exists(): + raise SystemExit(f"Missing {CACHE}. Build it first: python build_dataset.py") + d = np.load(CACHE) + X, y3, rat = d["X"], d["y"], d["rat"] + y = (y3 == 2).astype(np.int64) # 1 = pricking, 0 = movement + return X, y, rat + +def balanced(idx, y, rng): + i0, i1 = idx[y[idx] == 0], idx[y[idx] == 1] + n = min(len(i0), len(i1)) + return rng.permutation(np.concatenate([rng.choice(i0, n, False), rng.choice(i1, n, False)])) + +def stream_train(model, std, X, idx, y, feat=extract_features, update_std=True): + """Stream one sample at a time; prequential (test-then-train) accuracy curve.""" + correct, curve = 0, [] + for i, j in enumerate(idx, 1): + f = feat(X[j]) + if update_std: + std.update(f) + fz = std.transform(f) + correct += (model.predict(fz) == y[j]) # test first + model.update(fz, y[j]) # then train + if i % 50 == 0: + curve.append((i, correct / i)) + return curve + +def eval_acc(model, std, X, idx, y, feat=extract_features): + c = sum(model.predict(std.transform(feat(X[j]))) == y[j] for j in idx) + return c / len(idx) + + +# ── Demos ──────────────────────────────────────────────────────────────────── +def demo_online(X, y, rat, rng): + print("=" * 60) + print("DEMO 1 - online learning on RAT4") + print("=" * 60, flush=True) + idx = balanced(np.where(rat == 4)[0], y, rng) + cut = int(len(idx) * 0.7); tr, te = idx[:cut], idx[cut:] + model, std = LogRegSGD(N_FEAT), RunningStandardizer(N_FEAT) + curve = stream_train(model, std, X, tr, y) + for n, acc in curve[::max(1, len(curve) // 8)]: + print(f" after {n:>4} streamed samples: prequential acc = {acc:.1%}", flush=True) + print(f" final held-out test acc = {eval_acc(model, std, X, te, y):.1%}", flush=True) + +def demo_cross_rat(X, y, rat, rng): + print("\n" + "=" * 60) + print("DEMO 2 - cross-rat: train on RAT4/5/8/10, then adapt to RAT9") + print("=" * 60, flush=True) + src = balanced(np.where(np.isin(rat, [4, 5, 8, 10]))[0], y, rng) + model, std = LogRegSGD(N_FEAT), RunningStandardizer(N_FEAT) + stream_train(model, std, X, src, y) + tgt = balanced(np.where(rat == 9)[0], y, rng) + adapt, test = tgt[:len(tgt) // 2], tgt[len(tgt) // 2:] + print(f" before adapting, acc on RAT9 = {eval_acc(model, std, X, test, y):.1%}", flush=True) + model.lr = 0.02 + buf = [] + for blk in range(0, len(adapt), 100): + for j in adapt[blk:blk + 100]: + buf.append(j); model.update(std.transform(extract_features(X[j])), y[j]) + for _ in range(3): # small replay buffer + for j in rng.permutation(buf): + model.update(std.transform(extract_features(X[j])), y[j]) + print(f" after {len(buf):>4} RAT9 samples: acc = {eval_acc(model, std, X, test, y):.1%}", flush=True) + +def demo_channel_dropout(X, y, rat, rng): + print("\n" + "=" * 60) + print("DEMO 3 - channel dropout on RAT4") + print("=" * 60, flush=True) + idx = balanced(np.where(rat == 4)[0], y, rng) + cut = int(len(idx) * 0.6); tr, te = idx[:cut], idx[cut:] + model, std = LogRegSGD(N_FEAT), RunningStandardizer(N_FEAT) + stream_train(model, std, X, tr, y) + print(f" healthy electrode, RAT4 test acc = {eval_acc(model, std, X, te, y):.1%}", flush=True) + + dead = np.random.default_rng(7).choice(N_CH, int(0.4 * N_CH), False) + def feat_dropout(win): + w = win.copy(); w[dead] = 0.0 + return extract_features(w) + print(f" after {len(dead)} channels drop, test acc = " + f"{eval_acc(model, std, X, te, y, feat=feat_dropout):.1%}", flush=True) + print(" online recovery on dropped-channel stream (lr=0.02 + replay buffer):", flush=True) + model.lr = 0.02 + buf = [] + for blk in range(0, len(tr), 100): + for j in tr[blk:blk + 100]: + buf.append(j); model.update(std.transform(feat_dropout(X[j])), y[j]) + for _ in range(3): + for j in rng.permutation(buf): + model.update(std.transform(feat_dropout(X[j])), y[j]) + print(f" after {len(buf):>4} samples: acc = " + f"{eval_acc(model, std, X, te, y, feat=feat_dropout):.1%}", flush=True) + + +def main(): + t0 = time.time() + X, y, rat = load_task() + rng = np.random.default_rng(0) + demo_online(X, y, rat, rng) + demo_cross_rat(X, y, rat, rng) + demo_channel_dropout(X, y, rat, rng) + print(f"\nModel: {N_FEAT} weights + 1 bias = {N_FEAT + 1} parameters") + print(f"Time: {time.time() - t0:.0f}s", flush=True) + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/stream_raw.py b/scripts/cap_pipeline/ondevice/stream_raw.py new file mode 100644 index 0000000..9f853b3 --- /dev/null +++ b/scripts/cap_pipeline/ondevice/stream_raw.py @@ -0,0 +1,511 @@ +""" +stream_raw.py (on-device training + cross-rat adaptation demo) +-------------------------------------------------------------- +Replays raw 56-channel ENG to the Arduino sketch `ArduinoOnDeviceTrain`, which +detects CAPs live and trains a logistic-regression model on-device (pricking vs +movement). The board does its OWN detection; the laptop only chooses which CAPs +to send and in what order. + +The demo (M4 + M5): + 1. TRAIN — stream shuffled, class-balanced CAPs from several rats; the board + learns online (one SGD step per detected CAP). + 2. ADAPT — switch to a NEW rat (different cuff placement = distribution shift). + The model is mis-tuned for it, then recovers as it keeps training. + +To SEE the drop and recovery we can't use the board's running accuracy (online +SGD adapts so fast the dip hides in the noise). Instead, at checkpoints we stream +a FIXED held-out test set from the new rat with label 254 = "evaluate only": the +board predicts and reports but does not train or update its stats. The host scores +those predictions, giving a held-out accuracy at each checkpoint. + +Why CAPs are cut WITH their real preceding signal: that lead-in doubles as the +inter-event gap so the detector re-arms, and it reproduces the exact local noise +floor the EMA threshold was tuned for (synthesising gaps from random baseline +raises the floor and suppresses detection). + +Packet (119 bytes): [0xAA][0xBB] + label(1) + ts(4) + 56xuint16 LE(112). + label: 0=dorsi 1=plantar 2=pricking (board collapses to move/prick) ; + 254=eval-only ; 255=unknown (lead-in). +Replies: "CAP ts=N pred=P label=Y p=PROB acc=ACC" (a trained CAP) + "EVAL ts=N pred=P p=PROB" (a held-out prediction) + +Usage: + python scripts/cap_pipeline/ondevice/stream_raw.py --selftest # no hardware + python scripts/cap_pipeline/ondevice/stream_raw.py # stream to board + python scripts/cap_pipeline/ondevice/stream_raw.py --list-ports +""" + +import argparse +import csv +import pathlib +import struct +import sys +import threading +import time +from collections import deque + +import numpy as np +from scipy.signal import sosfilt, lfilter, lfilter_zi, butter, iirnotch, tf2sos + +HERE = pathlib.Path(__file__).resolve().parent +DATA_DIR = HERE.parents[2] / "data" +CACHE_DIR = HERE.parent / "caches" # git-ignored +sys.path.insert(0, str(HERE)) +from model import extract_features, RunningStandardizer, LogRegSGD, N_FEAT + +# ───────────────────────────────────────────────── +# CONFIGURATION +# ───────────────────────────────────────────────── + +SERIAL_PORT = "/dev/cu.usbmodem1101" +BAUD_RATE = 115200 +SLOWDOWN = 10 # replay at 1/SLOWDOWN real speed (10 -> 3 kHz effective) +BATCH_SZ = 60 # samples per serial write +SEED = 0 +OUT_CSV = CACHE_DIR / "adaptation_curve.csv" + +# Cross-rat scenario +TRAIN_RATS = [4, 5, 8, 10] # learn the task on these +ADAPT_RAT = 9 # then shift to this new rat and recover +CLASS_DIRS = ["DORSIFLEXION", "PLANTARFLEXION", "PRICKING"] # labels 0,1,2 + +MAX_TRAIN_EVENTS = 1200 # balanced CAPs streamed in the train phase +MAX_ADAPT_EVENTS = 800 # balanced CAPs from the new rat used for fine-tuning +EVAL_EVENTS = 120 # balanced held-out CAPs from the new rat (the test set) +EVAL_EVERY = 200 # adapt CAPs between held-out checkpoints + +# Stream geometry (samples @ 30 kHz). Each event is cut WITH its real lead-in. +LEAD = 2000 # real signal before the CAP (>> EMA time const ~200) +TRAIL = 300 # samples after the CAP +LABEL_PRE = 100 # CAP-vicinity that carries the real/eval label (rest = unknown) + +SYNC = bytes([0xAA, 0xBB]) +FS = 30_000.0 +LABEL_UNKNOWN, LABEL_EVAL = 255, 254 + +# ── Detection parameters (must match the board + build_dataset.py) ──────────── +N_CH, N_RINGS = 56, 7 +EMA_ALPHA, EMA_INIT = 0.005, 5.0 +THRESHOLD, REFRACTORY = 5.01, 150 +PEAK_UV = 15.0 +WIN_PRE, WIN_POST = 49, 51 # model window = [trig-49 .. trig+51) = 100 samples + +b_n, a_n = iirnotch(2000.0, 30.0, fs=FS) +SOS_NOTCH = tf2sos(b_n, a_n) +SOS_BP = butter(6, [1000.0, 3000.0], btype="band", fs=FS, output="sos") + + +# ───────────────────────────────────────────────── +# RHD READER (raw uint16) +# ───────────────────────────────────────────────── + +def _rqs(f): + r = f.read(4) + if len(r) < 4: + return "" + n = struct.unpack("= 1) or major > 1: n_temp = struct.unpack("= 3) or major > 1: f.read(2) + if major > 1: _rqs(f) + ng = struct.unpack(" 0 and en: + for _ in range(nch): + _rqs(f); _rqs(f) + _, _, st, ce, _, _ = struct.unpack(" keep first 56 + return None + spb, n_amp = hdr["spb"], hdr["n_amp"] + b_ts, b_amp = spb * 4, spb * n_amp * 2 + b_skip = ((spb // 4) * hdr["n_aux"] * 2 + hdr["n_supply"] * 2 + + hdr["n_temp"] * 2 + spb * hdr["n_adc"] * 2 + + (spb * 2 if hdr["n_dig_in"] > 0 else 0) + + (spb * 2 if hdr["n_dig_out"] > 0 else 0)) + blocks = [] + while True: + if len(f.read(b_ts)) < b_ts: break + amp = f.read(b_amp) + if len(amp) < b_amp: break + blocks.append(np.frombuffer(amp, dtype=np.uint16).reshape(n_amp, spb, order="F")) + if b_skip and len(f.read(b_skip)) < b_skip: break + except Exception: + return None + if not blocks: + return None + return np.concatenate(blocks, axis=1)[:N_CH].T.copy() # drop extras (RAT7 64->56) -> (N, 56) + + +# ───────────────────────────────────────────────── +# CAUSAL CAP DETECTION (same math as build_dataset.py / the board) +# ───────────────────────────────────────────────── + +def detect(raw_u16, warmup): + """raw uint16 (N,56) -> (trigs, uvr). trigs = detected peak indices.""" + uv = 0.195 * (raw_u16.astype(np.float64) - 32768.0) + ref = (uv[:, :8].mean(1) + uv[:, 48:].mean(1)) / 2.0 + uvr = uv - ref[:, None] + N = len(uvr) + rings = uvr[:, :N_RINGS * 8].reshape(N, N_RINGS, 8).mean(2) + for r in range(N_RINGS): + rings[:, r] = sosfilt(SOS_BP, sosfilt(SOS_NOTCH, rings[:, r])) + daa = np.zeros(N) + for r in range(N_RINGS): + if r == 0: daa += rings[:, r] + else: daa[r:] += rings[:-r, r] + zi = lfilter_zi([EMA_ALPHA], [1.0, -(1.0 - EMA_ALPHA)]) * EMA_INIT + ema, _ = lfilter([EMA_ALPHA], [1.0, -(1.0 - EMA_ALPHA)], np.abs(daa), zi=zi) + abs_daa = np.abs(daa) + cand = np.where(abs_daa[warmup:] > THRESHOLD * ema[warmup:])[0] + warmup + trigs, next_ok = [], warmup + for t in cand: + if t < next_ok: + continue + next_ok = t + REFRACTORY + if abs_daa[t] <= PEAK_UV: + trigs.append(int(t)) + return np.array(trigs, dtype=np.int64), uvr + + +# ───────────────────────────────────────────────── +# EVENT COLLECTION + STREAM ASSEMBLY +# ───────────────────────────────────────────────── + +def collect_events(rats): + """Detect CAPs for the given rats; return list of (raw_chunk (LEAD+TRAIL,56), cls). + + Each chunk is one CAP plus its real preceding signal; cls = 1 pricking, 0 move. + """ + events = [] + for rat in rats: + before = len(events) + for ci, cdir in enumerate(CLASS_DIRS): + d = DATA_DIR / f"RAT{rat}_{cdir}" + cls = 1 if ci == 2 else 0 + for p in (sorted(d.glob("*.rhd")) if d.exists() else []): + raw = read_rhd_raw(p) + if raw is None: + continue + N = len(raw) + trigs, _ = detect(raw, warmup=int(0.5 * FS)) + for t in trigs: + if t - LEAD >= 0 and t + TRAIL <= N: + events.append((raw[t - LEAD: t + TRAIL], cls)) + print(f" RAT{rat}: {len(events) - before} CAPs", flush=True) + return events + + +def balance(events, cap, rng): + """Balanced, shuffled subset with at most `cap` events total.""" + mov = [e for e in events if e[1] == 0] + pri = [e for e in events if e[1] == 1] + rng.shuffle(mov); rng.shuffle(pri) + k = min(len(mov), len(pri), cap // 2) + sel = mov[:k] + pri[:k] + rng.shuffle(sel) + return sel + + +def make_segment(events, mode): + """Assemble events into a contiguous (raw, label, truth) block. + + mode 'train' -> label = class (board trains); mode 'eval' -> label = 254 + (board predicts only). The lead-in is 'unknown' so a stray boundary trigger + isn't trained/counted. truth = cls in the CAP vicinity else -1. + """ + raw_parts, lab_parts, tru_parts = [], [], [] + for chunk, cls in events: + lab = np.full(len(chunk), LABEL_UNKNOWN, np.uint8) + tru = np.full(len(chunk), -1, np.int8) + a, b = LEAD - LABEL_PRE, LEAD + TRAIL + lab[a:b] = LABEL_EVAL if mode == "eval" else (2 if cls == 1 else 0) + tru[a:b] = cls + raw_parts.append(chunk); lab_parts.append(lab); tru_parts.append(tru) + return (np.concatenate(raw_parts), np.concatenate(lab_parts), + np.concatenate(tru_parts)) + + +def build_run(rng): + """Build the ordered list of stages for the whole demo. + + Returns list of dicts: {name, kind ('train'|'eval'), raw, label, truth, ckpt}. + """ + print(f"Detecting CAPs: train rats {TRAIN_RATS} …", flush=True) + train_ev = balance(collect_events(TRAIN_RATS), MAX_TRAIN_EVENTS, rng) + print(f" {len(train_ev)} balanced training CAPs", flush=True) + + print(f"Detecting CAPs: adapt rat {ADAPT_RAT} …", flush=True) + rat_ev = collect_events([ADAPT_RAT]) + rng.shuffle(rat_ev) + test_ev = balance(rat_ev, EVAL_EVENTS, rng) + test_ids = {id(e) for e in test_ev} + adapt_ev = balance([e for e in rat_ev if id(e) not in test_ids], MAX_ADAPT_EVENTS, rng) + print(f" {len(adapt_ev)} adapt CAPs, {len(test_ev)} held-out test CAPs", flush=True) + + def eval_stage(ckpt): + raw, lab, tru = make_segment(test_ev, "eval") + return dict(name=f"eval@{ckpt}", kind="eval", raw=raw, label=lab, truth=tru, ckpt=ckpt) + + def train_stage(name, ev): + raw, lab, tru = make_segment(ev, "train") + return dict(name=name, kind="train", raw=raw, label=lab, truth=tru, ckpt=-1) + + stages = [train_stage(f"train: rats {TRAIN_RATS}", train_ev)] + done = 0 + stages.append(eval_stage(done)) # before adapting + for blk in range(0, len(adapt_ev), EVAL_EVERY): + chunk = adapt_ev[blk:blk + EVAL_EVERY] + done += len(chunk) + stages.append(train_stage(f"adapt RAT{ADAPT_RAT} (+{len(chunk)})", chunk)) + stages.append(eval_stage(done)) + return stages + + +# ───────────────────────────────────────────────── +# OFFLINE SELF-TEST (no hardware — predicts what the board will do) +# ───────────────────────────────────────────────── + +def selftest(stages): + # Concatenate everything (the board's filters/EMA run continuously) and tag + # every sample with its stage so detections can be routed. + raws = np.concatenate([s["raw"] for s in stages]) + truth = np.concatenate([s["truth"] for s in stages]) + sid = np.concatenate([np.full(len(s["raw"]), i, np.int32) for i, s in enumerate(stages)]) + print(f"\nAssembled run: {len(raws)} samples ({len(raws)/FS:.1f} s, " + f"{len(raws)/FS*SLOWDOWN:.0f} s wall @ slowdown {SLOWDOWN})") + + trigs, uvr = detect(raws, warmup=0) + in_cap = truth[trigs] >= 0 + print(f"Detections: {len(trigs)} (in CAP {int(in_cap.sum())}, in lead-in {int((~in_cap).sum())})") + + model, std = LogRegSGD(N_FEAT), RunningStandardizer(N_FEAT) + recent = deque(maxlen=200) # windowed training accuracy + train_seen = 0 + checkpoints = {} # ckpt -> [correct, total] + cur = -1 + + def flush_eval(stage): + c, n = checkpoints.get(stage["ckpt"], [0, 0]) + acc = c / n if n else 0.0 + print(f" held-out acc = {acc:.3f} ({c}/{n}) {'█' * int(acc * 30)}", flush=True) + + for t in trigs: + if t - WIN_PRE < 0 or t + WIN_POST > len(uvr): + continue + y = int(truth[t]) + if y < 0: + continue + si = int(sid[t]); st = stages[si] + if si != cur: # stage boundary + if cur >= 0 and stages[cur]["kind"] == "eval": + flush_eval(stages[cur]) + cur = si + if st["kind"] == "train": + print(f"\n{st['name']}", flush=True) + else: + print(f"\n eval on RAT{ADAPT_RAT} (after {st['ckpt']} adapt CAPs):", flush=True) + + window = uvr[t - WIN_PRE: t + WIN_POST].T + f = extract_features(window) + if st["kind"] == "eval": + z = std.transform(f) # no std.update, no model.update + a = checkpoints.setdefault(st["ckpt"], [0, 0]) + a[0] += int(model.predict(z) == y); a[1] += 1 + else: + std.update(f) + z = std.transform(f) + recent.append(int(model.predict(z) == y)) + train_seen += 1 + model.update(z, y) + if train_seen % 200 == 0: + print(f" trained {train_seen:4d} CAPs, last-{len(recent)} acc {sum(recent)/len(recent):.3f}", + flush=True) + if cur >= 0 and stages[cur]["kind"] == "eval": + flush_eval(stages[cur]) + + print("\nHeld-out accuracy on the new rat, by checkpoint:") + for ck in sorted(checkpoints): + c, n = checkpoints[ck] + acc = c / n if n else 0.0 + print(f" after {ck:4d} adapt CAPs: {acc:.3f}", flush=True) + + +# ───────────────────────────────────────────────── +# ARDUINO REPLY READER (background thread) +# ───────────────────────────────────────────────── + +class ReplyReader(threading.Thread): + def __init__(self, ser, state, csv_writer): + super().__init__(daemon=True) + self.ser, self.state, self.csv = ser, state, csv_writer + self.train_total = self.train_correct = 0 + self.ckpt = {} # ckpt -> [correct, total] + + def run(self): + buf = b"" + while True: + try: + chunk = self.ser.read(64) + if not chunk: + continue + buf += chunk + while b"\n" in buf: + line, buf = buf.split(b"\n", 1) + self._handle(line.decode("ascii", errors="ignore").strip()) + except Exception: + break + + def _handle(self, line): + if line.startswith("EVAL"): + kv = {k: v for k, v in (p.split("=") for p in line.split() if "=" in p)} + pred = int(kv.get("pred", -1)) + ts = int(kv.get("ts", -1)) + truths = self.state["eval_truths"] # one per held-out event, in order + ev = ts // (LEAD + TRAIL) # which event this detection is in + scored = self.state["eval_scored"] + if 0 <= ev < len(truths) and ev not in scored: # ts-based: robust to misses/dupes + scored.add(ev) + a = self.ckpt.setdefault(self.state["ckpt"], [0, 0]) + a[0] += int(pred == truths[ev]); a[1] += 1 + return + if not line.startswith("CAP"): + if line: + print(f" [Arduino] {line}") + return + kv = {k: v for k, v in (p.split("=") for p in line.split() if "=" in p)} + pred, label = int(kv.get("pred", -1)), int(kv.get("label", -1)) + if label >= 0: + self.train_total += 1 + self.train_correct += int(pred == label) + if self.csv: + self.csv.writerow(["train", self.state["stage"], kv.get("ts", ""), + pred, label, kv.get("p", ""), kv.get("acc", ""), ""]) + + +# ───────────────────────────────────────────────── +# SERIAL STREAMING +# ───────────────────────────────────────────────── + +def stream_to_board(stages, port, write_csv): + import serial + state = {"stage": "-", "ckpt": -1, "eval_truths": [], "eval_scored": set()} + + csv_file = csv_writer = None + if write_csv: + CACHE_DIR.mkdir(parents=True, exist_ok=True) + csv_file = open(OUT_CSV, "w", newline="") + csv_writer = csv.writer(csv_file) + csv_writer.writerow(["kind", "stage", "ts", "pred", "label", "p", "board_acc", "held_out_acc"]) + + print(f"Opening {port} …") + ser = serial.Serial(port, BAUD_RATE, timeout=0.05) + time.sleep(2) + reader = ReplyReader(ser, state, csv_writer) + reader.start() + batch_dt = BATCH_SZ / (FS / SLOWDOWN) + + def send(raw, label): + batch = bytearray(); t0 = time.perf_counter() + for i in range(len(raw)): + batch += SYNC + struct.pack(" 0: + time.sleep(wait) + ser.write(bytes(batch)); batch = bytearray(); t0 = time.perf_counter() + if batch: + ser.write(bytes(batch)) + + try: + for st in stages: + state["stage"] = st["name"] + if st["kind"] == "eval": + # one held-out truth per event, indexed by ts // (LEAD+TRAIL) + state["ckpt"] = st["ckpt"] + state["eval_truths"] = [int(c) for c in st["truth"] if c >= 0][::(TRAIL + LABEL_PRE)] + state["eval_scored"] = set() + print(f"\n eval checkpoint @ {st['ckpt']} adapt CAPs …") + else: + print(f"\n{'═'*54}\n {st['name']}\n{'═'*54}") + send(st["raw"], st["label"]) + time.sleep(0.5) # let trailing replies arrive + if st["kind"] == "eval": + a = reader.ckpt.get(st["ckpt"], [0, 0]) + acc = a[0] / a[1] if a[1] else 0.0 + print(f" held-out acc = {acc:.3f} ({a[0]}/{a[1]})") + if csv_writer: + csv_writer.writerow(["eval", st["name"], "", "", "", "", "", f"{acc:.3f}"]) + except KeyboardInterrupt: + print("\nStopped by user.") + finally: + ser.close() + if csv_file: + csv_file.close() + print(f"\nCurve → {OUT_CSV}") + if reader.ckpt: + cks = sorted(reader.ckpt) + b0 = reader.ckpt[cks[0]]; b1 = reader.ckpt[cks[-1]] + print(f"Recovery: {b0[0]/max(1,b0[1]):.3f} -> {b1[0]/max(1,b1[1]):.3f}") + + +# ───────────────────────────────────────────────── +# MAIN +# ───────────────────────────────────────────────── + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--selftest", action="store_true", help="validate offline (no hardware)") + ap.add_argument("--list-ports", action="store_true") + ap.add_argument("--port", default=SERIAL_PORT) + ap.add_argument("--no-csv", action="store_true") + args = ap.parse_args() + + if args.list_ports: + import serial.tools.list_ports + for p in serial.tools.list_ports.comports(): + print(f" {p.device} {p.description}") + return + + rng = np.random.default_rng(SEED) + stages = build_run(rng) + if args.selftest: + selftest(stages) + else: + stream_to_board(stages, args.port, write_csv=not args.no_csv) + + +if __name__ == "__main__": + main() diff --git a/scripts/cap_pipeline/ondevice/task_data.py b/scripts/cap_pipeline/ondevice/task_data.py new file mode 100644 index 0000000..9b5fd4e --- /dev/null +++ b/scripts/cap_pipeline/ondevice/task_data.py @@ -0,0 +1,84 @@ +""" +task_data.py + +Shared data setup for the model comparison. Defines the task (pricking vs +movement) and the evaluation split once, so every model script trains and tests +on exactly the same data and the results are directly comparable. + +Split: within-rat, temporal — first 70% of each class (in detection order) for +training, last 30% for test, balanced. The temporal split stops adjacent windows +from leaking between train and test. + +Each model script saves its per-test-sample probabilities with save_preds(); the +ensemble loads them with load_preds(). +""" +import sys, pathlib +import numpy as np + +HERE = pathlib.Path(__file__).resolve().parent +sys.path.insert(0, str(HERE)) +from model import extract_features + +CACHE = HERE.parent / "caches" / "cap_windows_rich.npz" +PRED_DIR = HERE.parent / "caches" / "model_preds" +RATS = [4, 5, 8, 9, 10] # rats with both classes (RAT7 excluded: files are truncated) +SEED = 0 + + +def load(): + d = np.load(CACHE) + X, y3, rat = d["X"].astype(np.float32), d["y"], d["rat"] + y = (y3 == 2).astype(np.int64) # 1 = pricking, 0 = movement + return X, y3, rat, y + + +def splits(): + """List of (rat, train_idx, test_idx), balanced, deterministic.""" + _, y3, rat, _ = load() + rng = np.random.default_rng(SEED) + out = [] + for r in RATS: + tr, te = {0: [], 1: []}, {0: [], 1: []} + for c in (0, 1, 2): + idx = np.where((rat == r) & (y3 == c))[0] # already time-ordered + cut = int(len(idx) * 0.7) + b = 1 if c == 2 else 0 + tr[b].extend(idx[:cut]); te[b].extend(idx[cut:]) + + def bal(d): + n = min(len(d[0]), len(d[1])) + return rng.permutation(np.concatenate( + [rng.choice(d[0], n, False), rng.choice(d[1], n, False)])) + out.append((r, bal(tr), bal(te))) + return out + + +def feats(X, idx): + """113 hand features per window (what the linear model / RF use).""" + return np.stack([extract_features(X[j]) for j in idx]) + + +def save_preds(name, rats, probs, labels): + """Save per-test-sample P(pricking) so the ensemble can combine models.""" + PRED_DIR.mkdir(parents=True, exist_ok=True) + np.savez(PRED_DIR / f"{name}.npz", + rat=np.asarray(rats), prob=np.asarray(probs), label=np.asarray(labels)) + + +def load_preds(name): + d = np.load(PRED_DIR / f"{name}.npz") + return d["rat"], d["prob"], d["label"] + + +def report(name, rats, probs, labels): + """Print per-rat and mean accuracy for one model.""" + rats, probs, labels = np.asarray(rats), np.asarray(probs), np.asarray(labels) + pred = (probs >= 0.5).astype(int) + print(f"\n{name} — within-rat, temporal split:") + accs = [] + for r in RATS: + m = rats == r + a = (pred[m] == labels[m]).mean() + accs.append(a) + print(f" RAT{r:<2d}: {a:.3f} ({m.sum()} test)") + print(f" mean over {len(RATS)} rats: {np.mean(accs):.3f} ± {np.std(accs):.3f}")