From 0139723329fbc1ab811280ebc9c24274dc755d51 Mon Sep 17 00:00:00 2001
From: Kyle Johnsen <johnsenkyle13@gmail.com>
Date: Wed, 10 Apr 2024 15:55:13 -0400
Subject: [PATCH] fix Spiking.t_samp_ms

---
 cleo/ephys/spiking.py       |  3 ++-
 pyproject.toml              |  2 +-
 tests/ephys/test_spiking.py | 15 +++++++++------
 3 files changed, 12 insertions(+), 8 deletions(-)

diff --git a/cleo/ephys/spiking.py b/cleo/ephys/spiking.py
index 77d558e..e0c8054 100644
--- a/cleo/ephys/spiking.py
+++ b/cleo/ephys/spiking.py
@@ -65,7 +65,8 @@ def _update_saved_vars(self, t_ms, i, t_samp_ms):
         if self.probe.save_history:
             self.i = np.concatenate([self.i, i])
             self.t_ms = np.concatenate([self.t_ms, t_ms])
-            self.t_samp_ms = np.concatenate([self.t_samp_ms, [t_samp_ms]])
+            t_samp_ms_rep = np.full_like(t_ms, t_samp_ms)
+            self.t_samp_ms = np.concatenate([self.t_samp_ms, t_samp_ms_rep])
 
     def connect_to_neuron_group(
         self, neuron_group: NeuronGroup, **kwparams
diff --git a/pyproject.toml b/pyproject.toml
index 9a77cc0..271a08f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "cleosim"
-version = "0.14.1"
+version = "0.14.2"
 description = "Cleo: the Closed-Loop, Electrophysiology, and Optogenetics experiment simulation testbed"
 authors = [
     "Kyle Johnsen <kyle@kjohnsen.org>",
diff --git a/tests/ephys/test_spiking.py b/tests/ephys/test_spiking.py
index c8d5432..370bdd4 100644
--- a/tests/ephys/test_spiking.py
+++ b/tests/ephys/test_spiking.py
@@ -1,12 +1,10 @@
 """Tests for ephys.spiking module"""
-import pytest
 import numpy as np
-from brian2 import SpikeGeneratorGroup, ms, mm, Network
-import neo
 import quantities as pq
+from brian2 import Network, SpikeGeneratorGroup, mm, ms
 
 from cleo import CLSimulator
-from cleo.ephys import *
+from cleo.ephys import MultiUnitSpiking, Probe, SortedSpiking
 from cleo.ioproc import RecordOnlyProcessor
 
 
@@ -81,7 +79,7 @@ def test_MUS_multiple_contacts():
     assert np.sum(mus.i == 0) < len(indices)
     assert np.sum(mus.i == 1) < len(indices)
 
-    assert len(mus.i) == len(mus.t_ms)
+    assert len(mus.i) == len(mus.t_ms) == len(mus.t_samp_ms)
 
 
 def test_MUS_multiple_groups():
@@ -106,6 +104,8 @@ 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_ms) == len(mus.t_samp_ms)
+    assert np.all(mus.t_samp_ms == 10)
 
 
 def test_MUS_reset():
@@ -156,7 +156,10 @@ def test_SortedSpiking():
     assert all(i == [2, 3, 5])
 
     for i in (0, 1, 4):
-        assert not i in ss.i
+        assert i not in ss.i
+
+    assert ss.t_ms.shape == ss.i.shape == ss.t_samp_ms.shape
+    assert np.all(np.in1d(ss.t_samp_ms, [3, 4, 5, 6]))
 
 
 def _test_reset(spike_signal_class):