fixing bug in Ts.count

pynapple/core/time_series.py

@@ -1433,6 +1433,64 @@ def value_from(self, data, ep=None):
 
         return data.__class__(t, d, time_support=time_support, **kwargs)
 
+    def count(self, *args, **kwargs):
+        """
+        Count occurences of events within bin_size or within a set of bins defined as an IntervalSet.
+        You can call this function in multiple ways :
+
+        1. *tsd.count(bin_size=1, time_units = 'ms')*
+        -> Count occurence of events within a 1 ms bin defined on the time support of the object.
+
+        2. *tsd.count(1, ep=my_epochs)*
+        -> Count occurent of events within a 1 second bin defined on the IntervalSet my_epochs.
+
+        3. *tsd.count(ep=my_bins)*
+        -> Count occurent of events within each epoch of the intervalSet object my_bins
+
+        4. *tsd.count()*
+        -> Count occurent of events within each epoch of the time support.
+
+        bin_size should be seconds unless specified.
+        If bin_size is used and no epochs is passed, the data will be binned based on the time support of the object.
+
+        Parameters
+        ----------
+        bin_size : None or float, optional
+            The bin size (default is second)
+        ep : None or IntervalSet, optional
+            IntervalSet to restrict the operation
+        time_units : str, optional
+            Time units of bin size ('us', 'ms', 's' [default])
+
+        Returns
+        -------
+        out: Tsd
+            A Tsd object indexed by the center of the bins.
+
+        Examples
+        --------
+        This example shows how to count events within bins of 0.1 second.
+
+        >>> import pynapple as nap
+        >>> import numpy as np
+        >>> t = np.unique(np.sort(np.random.randint(0, 1000, 100)))
+        >>> ts = nap.Ts(t=t, time_units='s')
+        >>> bincount = ts.count(0.1)
+
+        An epoch can be specified:
+
+        >>> ep = nap.IntervalSet(start = 100, end = 800, time_units = 's')
+        >>> bincount = ts.count(0.1, ep=ep)
+
+        And bincount automatically inherit ep as time support:
+
+        >>> bincount.time_support
+        >>>    start    end
+        >>> 0  100.0  800.0
+        """
+        t, d, ep = super().count(*args, **kwargs)
+        return Tsd(t=t, d=d, time_support=ep)
+
     def fillna(self, value):
         """
         Similar to pandas fillna function.

tests/test_time_series.py

@@ -2,7 +2,7 @@
 # @Author: gviejo
 # @Date:   2022-04-01 09:57:55
 # @Last Modified by:   gviejo
-# @Last Modified time: 2024-02-20 22:54:15
+# @Last Modified time: 2024-02-29 10:34:07
 #!/usr/bin/env python
 
 """Tests of time series for `pynapple` package."""
@@ -1074,6 +1074,60 @@ def test_get(self, ts):
 
         assert len(ts[0:10,0]) == 10
 
+    def test_count(self, ts):
+        count = ts.count(1)
+        assert len(count) == 99
+        np.testing.assert_array_almost_equal(count.index, np.arange(0.5, 99, 1))
+
+        count = ts.count(bin_size=1)
+        assert len(count) == 99
+        np.testing.assert_array_almost_equal(count.index, np.arange(0.5, 99, 1))
+
+    def test_count_time_units(self, ts):
+        for b, tu in zip([1, 1e3, 1e6],['s', 'ms', 'us']):
+            count = ts.count(b, time_units = tu)
+            assert len(count) == 99
+            np.testing.assert_array_almost_equal(count.index, np.arange(0.5, 99, 1))
+
+            count = ts.count(b, tu)
+            assert len(count) == 99
+            np.testing.assert_array_almost_equal(count.index, np.arange(0.5, 99, 1))                
+
+    def test_count_with_ep(self, ts):
+        ep = nap.IntervalSet(start=0, end=100)
+        count = ts.count(1, ep)
+        assert len(count) == 100
+        np.testing.assert_array_almost_equal(count.values, np.ones(100))
+
+        count = ts.count(1, ep=ep)
+        assert len(count) == 100
+        np.testing.assert_array_almost_equal(count.values, np.ones(100))
+
+    def test_count_with_ep_only(self, ts):
+        ep = nap.IntervalSet(start=0, end=100)
+        count = ts.count(ep)
+        assert len(count) == 1
+        np.testing.assert_array_almost_equal(count.values, np.array([100]))
+
+        count = ts.count(ep=ep)
+        assert len(count) == 1
+        np.testing.assert_array_almost_equal(count.values, np.array([100]))
+
+        count = ts.count()
+        assert len(count) == 1
+        np.testing.assert_array_almost_equal(count.values, np.array([100]))
+
+
+    def test_count_errors(self, ts):        
+        with pytest.raises(ValueError):
+            ts.count(bin_size = {})
+
+        with pytest.raises(ValueError):
+            ts.count(ep = {})
+
+        with pytest.raises(ValueError):
+            ts.count(time_units = {})
+
 
 ####################################################
 # Test for tsdtensor