Add support for sending spike event & waveform from C++ to Python

anjaldoshi · anjaldoshi · commit 53f656869dcf · 2023-05-18T17:27:42.000-07:00
diff --git a/Modules/examples/bandpass_filter.py b/Modules/examples/bandpass_filter.py
@@ -70,25 +70,15 @@ def stop_acquisition(self):
         pass
     
     def handle_ttl_event(self, source_node, channel, sample_number, line, state):
-        """
-        Handle each incoming ttl event.
-        
-        Parameters:
-        source_node (int): id of the processor this event was generated from
-        channel (str): name of the event channel
-        sample_number (int): sample number of the event
-        line (int): the line on which event was generated (0-255) 
-        state (bool): event state true (ON) or false (OFF)
-        """
+        """ Handle each incoming ttl event """
         pass
 
-    def start_recording(self, recording_dir):
-        """ 
-        Called when recording starts
+    def handle_spike(self, source_node, electrode_name, num_channels, num_samples, sample_number, sorted_id, spike_data):
+        """ Handle each incoming spike """
+        pass
 
-        Parameters:
-        recording_dir - recording directory to be used by future record nodes.
-        """
+    def start_recording(self, recording_dir):
+        """ Called when recording starts """
         pass
     
     def stop_recording(self):
diff --git a/Modules/template/processor_template.py b/Modules/template/processor_template.py
@@ -21,7 +21,7 @@ def process(self, data):
         Process each incoming data buffer.
         
         Parameters:
-        data - numpy array.
+        data - N x M numpy array, where N = num_channles, M = num of samples in the buffer.
         """
         try:
             pass
@@ -48,6 +48,22 @@ def handle_ttl_event(self, source_node, channel, sample_number, line, state):
         state (bool): event state True (ON) or False (OFF)
         """
         pass
+
+    def handle_spike(self, source_node, electrode_name, num_channels, num_samples, sample_number, sorted_id, spike_data):
+        """
+        Handle each incoming spike.
+        
+        Parameters:
+        source_node (int): id of the processor this spike was generated from
+        electrode_name (str): name of the electrode
+        num_channels (int): number of channels associated with the electrode type
+        num_samples (int): total number of samples in the spike waveform 
+        sample_number (int): sample number of the spike
+        sorted_id (int): the sorted ID for this spike
+        spike_data (numpy array): N x M numpy array, where N = num_channels & M = num_samples (read-only).
+        """
+        # print("SPIKE RECEIVED! Source Node:", source_node, ", Electrode name:", electrode_name, ", Num channels:", num_channels, ", Sample num:", sample_number, ", Sorted ID:", sorted_id)
+        pass
     
     def start_recording(self, recording_dir):
         """ 
diff --git a/Source/PythonProcessor.cpp b/Source/PythonProcessor.cpp
@@ -25,7 +25,6 @@ along with this program.  If not, see <http://www.gnu.org/licenses/>.
 #include <filesystem>
 
 #include "PythonProcessor.h"
-#include "PythonProcessorEditor.h"
 
 namespace py = pybind11;
 
@@ -38,8 +37,8 @@ PYBIND11_EMBEDDED_MODULE(oe_pyprocessor, module){
 PythonProcessor::PythonProcessor()
     : GenericProcessor("Python Processor")
 {
-    pyModule = NULL;
-    pyObject = NULL;
+    pyModule = nullptr;
+    pyObject = nullptr;
     moduleReady = false;
     scriptPath = "";
     moduleName = "";
@@ -58,8 +57,6 @@ PythonProcessor::~PythonProcessor()
     if(Py_IsInitialized() > 0)
     {
         {
-            delete pyModule;
-            delete pyObject;
             py::gil_scoped_release release;
         }
         py::finalize_interpreter();
@@ -203,16 +200,41 @@ void PythonProcessor::handleTTLEvent(TTLEventPtr event)
 }
 
 
-// void PythonProcessor::handleSpike(SpikePtr event)
-// {
-//     // py::gil_scoped_acquire acquire;
-//     try {
-//         pyObject->attr("handle_spike_event")();
-//     }
-//     catch (py::error_already_set& e) {
-//         handlePythonException(e);
-//     }
-// }
+void PythonProcessor::handleSpike(SpikePtr spike)
+{
+    if (spike->getStreamId() == currentStream)
+    {
+        auto spikeChanInfo = spike->getChannelInfo();
+
+        const int sourceNodeId = spikeChanInfo->getSourceNodeId();
+        const String electrodeName = spikeChanInfo->getName();
+        const int numChans = spikeChanInfo->getNumChannels();
+        const int64 sampleNum = spike->getSampleNumber();
+        const uint16 sortedId = spike->getSortedId();
+        const int numSamples = spikeChanInfo->getTotalSamples();
+
+        py::array_t<float> spikeData = py::array_t<float>({ numChans, numSamples });
+
+        for (int i = 0; i < numChans; ++i) 
+        {
+            const float* spikeChanDataPtr = spike->getDataPointer(i);
+            float* numpyChannelPtr = spikeData.mutable_data(i, 0);
+            memcpy(numpyChannelPtr, spikeChanDataPtr, sizeof(float) * numSamples);
+        }
+
+        // py::gil_scoped_acquire acquire;
+        if(py::hasattr(*pyObject, "handle_spike"))
+        {
+            try {
+                pyObject->attr("handle_spike")
+                    (sourceNodeId, electrodeName.toRawUTF8(), numChans, numSamples, sampleNum, sortedId, spikeData);
+            }
+            catch (py::error_already_set& e) {
+                handlePythonException(e);
+            }
+        }
+    }
+}
 
 
 // void PythonProcessor::handleBroadcastMessage(String message)
@@ -249,52 +271,67 @@ bool PythonProcessor::startAcquisition()
     {
         // py::gil_scoped_acquire acquire;
 
-        try {
-            pyObject->attr("start_acquisition")();
-        }
-        catch (py::error_already_set& e) {
-            handlePythonException(e);
+        if(py::hasattr(*pyObject, "start_acquisition"))
+        {
+            try {
+                pyObject->attr("start_acquisition")();
+            }
+            catch (py::error_already_set& e) {
+                handlePythonException(e);
+            }
         }
         return true;
     }
     return false;
 }
 
-bool PythonProcessor::stopAcquisition() {
+bool PythonProcessor::stopAcquisition()
+{
     if (moduleReady)
     {
         // py::gil_scoped_acquire acquire;
-        try {
-            pyObject->attr("stop_acquisition")();
-        }
-        catch (py::error_already_set& e) {
-            handlePythonException(e);
+        if(py::hasattr(*pyObject, "stop_acquisition"))
+        {
+            try {
+                pyObject->attr("stop_acquisition")();
+            }
+            catch (py::error_already_set& e) {
+                handlePythonException(e);
+            }
         }
     }
     return true;
 }
 
-void PythonProcessor::startRecording() {
+void PythonProcessor::startRecording() 
+{
     String recordingDirectory = CoreServices::getRecordingDirectoryName();
 
     // py::gil_scoped_acquire acquire;
-    try {
-        pyObject->attr("start_recording")(recordingDirectory.toRawUTF8());
-    }
-    catch (py::error_already_set& e) {
-        handlePythonException(e);
+    if(moduleReady && py::hasattr(*pyObject, "start_recording"))
+    {
+        try {
+            pyObject->attr("start_recording")(recordingDirectory.toRawUTF8());
+        }
+        catch (py::error_already_set& e) {
+            handlePythonException(e);
+        }
     }
 }
 
 
 
-void PythonProcessor::stopRecording() {
+void PythonProcessor::stopRecording() 
+{
     // py::gil_scoped_acquire acquire;
-    try {
-        pyObject->attr("stop_recording")();
-    }
-    catch (py::error_already_set& e) {
-        handlePythonException(e);
+    if(moduleReady && py::hasattr(*pyObject, "stop_recording"))
+    {
+        try {
+            pyObject->attr("stop_recording")();
+        }
+        catch (py::error_already_set& e) {
+            handlePythonException(e);
+        }
     }
 }
 
diff --git a/Source/PythonProcessor.h b/Source/PythonProcessor.h
@@ -106,10 +106,10 @@ class PythonProcessor : public GenericProcessor
 	// 	the plugin's process() method */
 	void handleTTLEvent(TTLEventPtr event) override;
 
-	// /** Handles spikes received by the processor
-	// 	Called automatically for each received spike whenever checkForEvents(true) is called from
-	// 	the plugin's process() method */
-	// void handleSpike(SpikePtr spike) override;
+	/** Handles spikes received by the processor
+		Called automatically for each received spike whenever checkForEvents(true) is called from
+		the plugin's process() method */
+	void handleSpike(SpikePtr spike) override;
 
 	// /** Handles broadcast messages sent during acquisition
 	// 	Called automatically whenever a broadcast message is sent through the signal chain */
