fix spike count problem

cleo/ephys/spiking.py

@@ -12,11 +12,11 @@
 from attrs import define, field, fields
 from bidict import bidict
 from brian2 import NeuronGroup, Quantity, SpikeMonitor, mm, ms
-from nptyping import NDArray, Shape, UInt, Float
+from nptyping import Float, NDArray, Shape, UInt
 
 from cleo.base import NeoExportable
 from cleo.ephys.probes import Signal
-from cleo.utilities import unit_safe_cat
+from cleo.utilities import rng, unit_safe_cat
 
 
 @define(eq=False)
@@ -232,7 +232,7 @@ def _noisily_detect_spikes_per_channel(
         self, i_probe, t
     ) -> Tuple[NDArray, NDArray, NDArray]:
         probs_for_spikes = self._dtct_prob_array[i_probe]
-        detected_spikes = np.random.random(probs_for_spikes.shape) < probs_for_spikes
+        detected_spikes = rng.random(probs_for_spikes.shape) < probs_for_spikes
         y = np.sum(detected_spikes, axis=0)
         # ⬇ nonzero gives row, column indices of each nonzero element
         i_spike_detected, i_c_detected = detected_spikes.nonzero()
@@ -294,15 +294,18 @@ def get_state(
         # inherit docstring
         i_probe, t = self._get_new_spikes()
         i_probe, t = self._noisily_detect_spikes(i_probe, t)
-        y = np.zeros(len(self.i_probe_by_i_ng), dtype=bool)
-        y[i_probe] = 1
+        y = np.bincount(i_probe.astype(int))
+        # include 0s for upper indices not seen:
+        y = np.concatenate([y, np.zeros(len(self.i_probe_by_i_ng) - len(y))])
         t_samp = self.probe.sim.network.t / ms
         self._update_saved_vars(t, i_probe, t_samp)
         return (i_probe, t, y)
 
     def _noisily_detect_spikes(self, i_probe, t) -> Tuple[NDArray, NDArray]:
+        # dtct_prob_array: n_nrns x n_channels
         probs_for_spikes = self._dtct_prob_array[i_probe]
-        detected_spikes = np.random.random(probs_for_spikes.shape) < probs_for_spikes
+        # now n_spks x n_channels
+        detected_spikes = rng.random(probs_for_spikes.shape) < probs_for_spikes
         i_spike_detected = detected_spikes.nonzero()
         i_probe_out = i_probe[i_spike_detected]
         t_out = t[i_spike_detected]

tests/ephys/test_spiking.py

@@ -45,23 +45,26 @@ def test_MUS_multiple_contacts():
 
     # remember i here is channel, no longer neuron
     i, t, y = mus.get_state()
-    assert len(i) == 0
-    assert len(t) == 0
+    assert len(i) == len(t) == y.sum() == 0
 
     sim.run(1 * ms)  # 1 ms
     i, t, y = mus.get_state()
+    assert len(i) == len(t) == y.sum()
     assert (0 in i) and (0.9 * ms in t)
 
     sim.run(1 * ms)  # 2 ms
     i, t, y = mus.get_state()
+    assert len(i) == len(t) == y.sum()
     assert len(i) == 0 and len(t) == 0
 
     sim.run(1 * ms)  # 3 ms
     i, t, y = mus.get_state()
+    assert len(i) == len(t) == y.sum()
     assert (0 in i) and (2.1 * ms in t)
 
     sim.run(1 * ms)  # 4 ms
     i, t, y = mus.get_state()
+    assert len(i) == len(t) == y.sum()
     # should pick up 2 spikes on first and 1 or 2 on second channel
     assert 1 in i and len(i) >= 3
     assert 3.5 * ms in t
@@ -70,6 +73,7 @@ def test_MUS_multiple_contacts():
     sim.run(2 * ms)
     # sim.get_state() nested dict is what will be passed to IO processor
     i, t, y = sim.get_state()["Probe"]["mus"]
+    assert len(i) == len(t) == y.sum()
     assert (0 in i) and (1 in i) and len(i) >= 7
     assert all(t_i in (t / ms).round(2) for t_i in [4.1, 4.9, 5.1, 5.3, 5.5])
 
@@ -104,7 +108,7 @@ def test_MUS_multiple_groups():
     assert 20 < np.sum(mus.i == 0) < 60
     # second channel would have caught all spikes from sgg1 and sgg2
     assert np.sum(mus.i == 1) == 60
-    assert len(mus.t) == len(mus.t_samp)
+    assert len(mus.t) == len(mus.t_samp) == y.sum()
     assert np.all(mus.t_samp == 10 * ms)
 
 
@@ -142,18 +146,22 @@ def test_SortedSpiking():
     sim.run(3 * ms)  # 3 ms
     i, t, y = sim.get_state()["Probe"]["ss"]
     assert all(i == [2, 2])
+    assert len(i) == len(t) == y.sum()
 
     sim.run(1 * ms)  # 4 ms
     i, t, y = ss.get_state()
     assert all(i == [2, 3])
+    assert len(i) == len(t) == y.sum()
 
     sim.run(1 * ms)  # 5 ms
     i, t, y = ss.get_state()
     assert all(i == [3, 5])
+    assert len(i) == len(t) == y.sum()
 
     sim.run(1 * ms)  # 6 ms
     i, t, y = ss.get_state()
     assert all(i == [2, 3, 5])
+    assert len(i) == len(t) == y.sum()
 
     for i in (0, 1, 4):
         assert i not in ss.i