callback for returning spikes from CORENEURON

src/nrniv/netpar.cpp

@@ -26,6 +26,8 @@ implementNrnHash(Gid2PreSyn, int, PreSyn*)
 #include <netcon.h>
 #include <cvodeobj.h>
 #include <netcvode.h>
+#include <vector>
+#include "ivocvect.h"
 
 #define BGP_INTERVAL 2
 #if BGP_INTERVAL == 2
@@ -35,6 +37,8 @@ static int n_bgp_interval;
 static Symbol* netcon_sym_;
 static Gid2PreSyn* gid2out_;
 static Gid2PreSyn* gid2in_;
+static IvocVect* all_spiketvec = NULL;
+static IvocVect* all_spikegidvec = NULL;
 static double t_exchange_;
 static double dt1_; // 1/dt
 static void alloc_space();
@@ -65,6 +69,15 @@ extern PreSyn* nrn_gid2presyn(int);
 extern int nrn_gid_exists(int);
 extern double nrnmpi_step_wait_; // barrier at beginning of spike exchange.
 
+/**
+ * @brief NEURON callback used from CORENEURON to transfer all spike vectors after simulation
+ *
+ * @param spiketvec - vector of spikes
+ * @param spikegidvec - vector of gids
+ * @return 1 if CORENEURON shall drop writing `out.dat`; 0 otherwise
+ */
+int nrnthread_all_spike_vectors_return(std::vector<double>& spiketvec, std::vector<int>& spikegidvec);
+
 // BGPDMA can be 0,1,2,3,6,7
 // (BGPDMA & 1) > 0 means multisend ISend allowed
 // (BGPDMA & 2) > 0 means multisend Blue Gene/P DMA allowed
@@ -1095,20 +1108,23 @@ void BBS::outputcell(int gid) {
 void BBS::spike_record(int gid, IvocVect* spikevec, IvocVect* gidvec) {
 	PreSyn* ps;
     if (gid >= 0) {
-	nrn_assert(gid2out_->find(gid, ps));
-	assert(ps);
-	ps->record(spikevec, gidvec, gid);
-    }else{ // record all output spikes
-	NrnHashIterate(Gid2PreSyn, gid2out_, PreSyn*, ps) {
-		if (ps->output_index_ >= 0) {
-			ps->record(spikevec, gidvec, ps->output_index_);
-		}
-	}}}
+        all_spiketvec = NULL, all_spikegidvec = NULL; // invalidate global spike vectors
+        nrn_assert(gid2out_->find(gid, ps));
+        assert(ps);
+        ps->record(spikevec, gidvec, gid);
+    }else{ // record all output spikes.
+        all_spiketvec = spikevec, all_spikegidvec = gidvec; // track global spike vectors (optimisation)
+        NrnHashIterate(Gid2PreSyn, gid2out_, PreSyn*, ps) {
+            if (ps->output_index_ >= 0) {
+                ps->record(all_spiketvec, all_spikegidvec, ps->output_index_);
+            }
+        }}}
     }
 }
 
 void BBS::spike_record(IvocVect* gids, IvocVect* spikevec, IvocVect* gidvec) {
 	int sz = vector_capacity(gids);
+    all_spiketvec = NULL, all_spikegidvec = NULL; // invalidate global spike vectors
 	double* pd = vector_vec(gids);
 	for (int i = 0; i < sz; ++i) {
 		PreSyn* ps;
@@ -1285,6 +1301,35 @@ void BBS::netpar_solve(double tstop) {
 	tstopunset;
 }
 
+int nrnthread_all_spike_vectors_return(std::vector<double>& spiketvec, std::vector<int>& spikegidvec){
+    assert(spiketvec.size() == spikegidvec.size());
+    if( spiketvec.size() ) {
+        /* Optimisation:
+         *  if all_spiketvec and all_spikegidvec are set (pc.spike_record with gid=-1 was called)
+         *  and they are still reference counted (obj_->refcount), then copy over incoming vectors.
+         */
+        if (all_spiketvec != NULL && all_spiketvec->obj_ != NULL  && all_spiketvec->obj_->refcount > 0 &&
+            all_spikegidvec != NULL && all_spikegidvec->obj_ != NULL && all_spikegidvec->obj_->refcount > 0) {
+
+            all_spiketvec->resize(spiketvec.size());
+            all_spikegidvec->resize(spikegidvec.size());
+            for (int i = 0; i < all_spiketvec->capacity(); ++i) {
+                all_spiketvec->elem(i) = spiketvec[i];
+                all_spikegidvec->elem(i) = spikegidvec[i];
+            }
+        }else{ // different underlying vectors for PreSyns
+            for (int i = 0; i < spikegidvec.size(); ++i ) {
+                PreSyn* ps;
+                if (gid2out_->find(spikegidvec[i], ps)) {
+                    ps->record(spiketvec[i]);
+                }
+            }
+        }
+    }
+    return 1;
+}
+
+
 static double set_mindelay(double maxdelay) {
 	double mindelay = maxdelay;
 	last_maxstep_arg_ = maxdelay;

src/nrniv/nrnbbcore_write.cpp

@@ -1868,6 +1868,10 @@ void nrnthread_trajectory_values(int tid, int n_pr, void** vpr, double t);
 void nrnthread_trajectory_return(int tid, int n_pr, int vecsz, void** vpr, double t);
 }
 
+extern "C" {
+extern int nrnthread_all_spike_vectors_return(std::vector<double>& spiketvec, std::vector<int>& spikegidvec);
+}
+
 static core2nrn_callback_t cnbs[]  = {
   {"nrn2core_group_ids_", (CNB)nrnthread_group_ids},
   {"nrn2core_mkmech_info_", (CNB)write_memb_mech_types_direct},
@@ -1888,6 +1892,8 @@ static core2nrn_callback_t cnbs[]  = {
   {"nrn2core_get_trajectory_requests_", (CNB)nrnthread_get_trajectory_requests},
   {"nrn2core_trajectory_values_", (CNB)nrnthread_trajectory_values},
   {"nrn2core_trajectory_return_", (CNB)nrnthread_trajectory_return},
+
+  {"nrn2core_all_spike_vectors_return_", (CNB)nrnthread_all_spike_vectors_return},
   {NULL, NULL}
 };
 

test/CMakeLists.txt

@@ -122,7 +122,20 @@ if(NRN_ENABLE_PYTHON AND PYTEST_FOUND)
     add_test(NAME coreneuron_direct_hoc COMMAND ${CMAKE_BINARY_DIR}/bin/nrniv
                                             ${PROJECT_SOURCE_DIR}/test/coreneuron/test_direct.hoc
                                     WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/test/mod)
-    list(APPEND TESTS coreneuron_direct_py coreneuron_direct_hoc)
+    add_test(NAME coreneuron_spikes_py COMMAND ${PYTHON_EXECUTABLE}
+                                            ${PROJECT_SOURCE_DIR}/test/coreneuron/test_spikes.py
+                                    WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/test/mod)
+    list(APPEND TESTS coreneuron_direct_py coreneuron_direct_hoc coreneuron_spikes_py)
+    if(NRN_ENABLE_MPI)
+      find_python_module(mpi4py)
+      if(mpi4py_FOUND)
+        add_test(
+                NAME coreneuron_mpi_spikes_py  COMMAND ${MPIEXEC} -np 2 ${PYTHON_EXECUTABLE}
+                                              ${PROJECT_SOURCE_DIR}/test/coreneuron/test_spikes.py mpi4py
+                                      WORKING_DIRECTORY ${PROJECT_BINARY_DIR}/test/mod)
+        list(APPEND TESTS coreneuron_mpi_spikes_py)
+      endif()
+    endif()
   endif()
 endif()
 set_tests_properties(${TESTS} PROPERTIES ENVIRONMENT "${TEST_ENV}")

test/coreneuron/test_spikes.py

@@ -0,0 +1,72 @@
+import pytest
+import sys
+
+from neuron import h, gui
+
+def test_spikes(use_mpi4py=False):
+    if use_mpi4py:
+        from mpi4py import MPI
+    h('''create soma''')
+    h.soma.L=5.6419
+    h.soma.diam=5.6419
+    h.soma.insert("hh")
+    h.soma.nseg = 3
+    ic = h.IClamp(h.soma(.25))
+    ic.delay = .1
+    ic.dur = 0.1
+    ic.amp = 0.3
+
+    ic2 = h.IClamp(h.soma(.75))
+    ic2.delay = 5.5
+    ic2.dur = 1
+    ic2.amp = 0.3
+
+    h.tstop = 10
+    h.cvode.use_fast_imem(1)
+    h.cvode.cache_efficient(1)
+
+    pc = h.ParallelContext()
+
+    pc.set_gid2node(pc.id()+1, pc.id())
+    myobj = h.NetCon(h.soma(0.5)._ref_v, None, sec=h.soma)
+    pc.cell(pc.id()+1, myobj)
+
+
+    # NEURON spikes run
+    nrn_spike_t = h.Vector()
+    nrn_spike_gids = h.Vector()
+    pc.spike_record(-1, nrn_spike_t, nrn_spike_gids)
+
+    h.run()
+
+    nrn_spike_t = nrn_spike_t.to_python()
+    nrn_spike_gids = nrn_spike_gids.to_python()
+
+    # CORENEURON spike_record(-1) / spike_record(gidlist):
+    from neuron import coreneuron
+    coreneuron.enable = True
+    coreneuron.verbose = 0
+    h.stdinit()
+    corenrn_all_spike_t = h.Vector()
+    corenrn_all_spike_gids = h.Vector()
+
+    pc.spike_record( -1 if pc.id() == 0 else (pc.id()),
+                     corenrn_all_spike_t,
+                     corenrn_all_spike_gids )
+    pc.psolve(h.tstop)
+
+    corenrn_all_spike_t = corenrn_all_spike_t.to_python()
+    corenrn_all_spike_gids = corenrn_all_spike_gids.to_python()
+
+
+    # check spikes match
+    assert(len(nrn_spike_t)) # check we've actually got spikes
+    assert(len(nrn_spike_t) == len(nrn_spike_gids)); # matching no. of gids
+    assert(nrn_spike_t == corenrn_all_spike_t)
+    assert(nrn_spike_gids == corenrn_all_spike_gids)
+
+    h.quit()
+
+
+if __name__ == "__main__":
+    test_spikes('mpi4py' in sys.argv)