Cleaning perievent

Author: gviejo

pynapple/process/perievent.py

@@ -2,15 +2,15 @@
 Functions to realign time series relative to a reference time.
 """
 
+import inspect
+from functools import wraps
+from numbers import Number
+
 import numpy as np
 
 from .. import core as nap
 from ._process_functions import _perievent_continuous, _perievent_trigger_average
 
-import inspect
-from functools import wraps
-from numbers import Number
-
 
 def _validate_perievent_inputs(func):
     @wraps(func)
@@ -23,18 +23,18 @@ def wrapper(*args, **kwargs):
             "timestamps": (nap.Ts, nap.Tsd, nap.TsdFrame, nap.TsdTensor, nap.TsGroup),
             "timeseries": (nap.Tsd, nap.TsdFrame, nap.TsdTensor),
             "tref": (nap.Ts, nap.Tsd, nap.TsdFrame, nap.TsdTensor),
-            "group": nap.TsGroup,
-            "ep": (nap.IntervalSet, None),
+            "group": (nap.TsGroup,),
+            "ep": (nap.IntervalSet,),
             "feature": (nap.Tsd, nap.TsdFrame, nap.TsdTensor),
-            "binsize": Number,
+            "binsize": (Number,),
             "windowsize": (tuple, Number),
-            "time_units": str,
+            "time_unit": (str,),
         }
         for param, param_type in parameters_type.items():
             if param in kwargs:
                 if not isinstance(kwargs[param], param_type):
                     raise TypeError(
-                        f"Invalid type. Parameter {param} must be of type {param_type}."
+                        f"Invalid type. Parameter {param} must be of type {[p.__name__ for p in param_type]}."
                     )
 
         # Call the original function with validated inputs
@@ -43,7 +43,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def _align_tsd(tsd, tref, window, time_support):
+def _align_tsd(tsd, tref, window, new_time_support):
     """
     Helper function compiled with numba for aligning times.
     See compute_perievent for using this function
@@ -72,21 +72,21 @@ def _align_tsd(tsd, tref, window, time_support):
     if isinstance(tsd, nap.Ts):
         for i in range(len(tref)):
             tmp = tsd.index[lbounds[i] : rbounds[i]] - tref.index[i]
-            group[i] = nap.Ts(t=tmp, time_support=time_support)
+            group[i] = nap.Ts(t=tmp, time_support=new_time_support)
     else:
         for i in range(len(tref)):
             tmp = tsd.index[lbounds[i] : rbounds[i]] - tref.index[i]
             tmp2 = tsd.values[lbounds[i] : rbounds[i]]
-            group[i] = nap.Tsd(t=tmp, d=tmp2, time_support=time_support)
+            group[i] = nap.Tsd(t=tmp, d=tmp2, time_support=new_time_support)
 
-    group = nap.TsGroup(group, time_support=time_support, bypass_check=True)
+    group = nap.TsGroup(group, time_support=new_time_support, bypass_check=True)
     group.set_info(ref_times=tref.index)
 
     return group
 
 
 @_validate_perievent_inputs
-def compute_perievent(timestamps, tref, windowsize, ep=None, time_unit="s"):
+def compute_perievent(timestamps, tref, windowsize, time_unit="s"):
     """
     Center the timestamps of a time series object or a time series group around the timestamps given by the `tref` argument.
     `windowsize` indicates the start and end of the window. If `windowsize=(-5, 10)`, the window will be from -5 second to 10 second.
@@ -104,8 +104,6 @@ def compute_perievent(timestamps, tref, windowsize, ep=None, time_unit="s"):
         The time reference of the event to align to
     windowsize : tuple of int/float or int or float
         The window size. Can be unequal on each side i.e. (-500, 1000).
-    ep : IntervalSet, optional
-        The epochs to perform the operation. If None, the default is the time support of the `timestamps` object.
     time_unit : str, optional
         Time units of the windowsize ('s' [default], 'ms', 'us').
 
@@ -137,20 +135,20 @@ def compute_perievent(timestamps, tref, windowsize, ep=None, time_unit="s"):
             "windowsize should be a tuple of 2 numbers or a single number."
         )
 
-    if ep is None:
-        ep = timestamps.time_support
-
     window = np.abs(nap.TsIndex.format_timestamps(np.array(windowsize), time_unit))
 
+    new_time_support = nap.IntervalSet(start=-window[0], end=window[1])
+
     if isinstance(timestamps, nap.TsGroup):
         toreturn = {}
         for n in timestamps.index:
-            toreturn[n] = _align_tsd(timestamps[n], tref, window, ep)
+            toreturn[n] = _align_tsd(timestamps[n], tref, window, new_time_support)
         return toreturn
     else:
-        return _align_tsd(timestamps, tref, window, ep)
+        return _align_tsd(timestamps, tref, window, new_time_support)
 
 
+@_validate_perievent_inputs
 def compute_perievent_continuous(timeseries, tref, windowsize, ep=None, time_unit="s"):
     """
     Center continuous time series around the timestamps given by the 'tref' argument.
@@ -230,6 +228,7 @@ def compute_perievent_continuous(timeseries, tref, windowsize, ep=None, time_uni
         return nap.TsdTensor(t=time_idx, d=new_data_array, time_support=time_support)
 
 
+@_validate_perievent_inputs
 def compute_event_trigger_average(
     group,
     feature,

tests/test_perievent.py

@@ -54,26 +54,52 @@ def test_compute_perievent_windowsize():
 def test_compute_perievent_raise_error():
     tsd = nap.Ts(t=np.arange(100))
     tref = nap.Ts(t=np.arange(10, 100, 10))
-    with pytest.raises(AssertionError) as e_info:
-        nap.compute_perievent(tsd, [0, 1, 2], windowsize=(-10, 10))
-    assert str(e_info.value) == "tref should be a Ts or Tsd object."
+    with pytest.raises(TypeError) as e_info:
+        nap.compute_perievent(tsd, tref=[0, 1, 2], windowsize=(-10, 10))
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter tref must be of type ['Ts', 'Tsd', 'TsdFrame', 'TsdTensor']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
-        nap.compute_perievent([0, 1, 2], tref, windowsize=(-10, 10))
-    assert str(e_info.value) == "data should be a Ts, Tsd or TsGroup."
+    with pytest.raises(TypeError) as e_info:
+        nap.compute_perievent(timestamps=[0, 1, 2], tref=tref, windowsize=(-10, 10))
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter timestamps must be of type ['Ts', 'Tsd', 'TsdFrame', 'TsdTensor', 'TsGroup']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_perievent(tsd, tref, windowsize={0: 1})
-    assert str(e_info.value) == "windowsize should be a tuple or int or float."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter windowsize must be of type ['tuple', 'Number']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_perievent(tsd, tref, windowsize=10, time_unit=1)
-    assert str(e_info.value) == "time_unit should be a str."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter time_unit must be of type ['str']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(RuntimeError) as e_info:
         nap.compute_perievent(tsd, tref, windowsize=10, time_unit="a")
     assert str(e_info.value) == "time_unit should be 's', 'ms' or 'us'"
 
+    with pytest.raises(RuntimeError) as e_info:
+        nap.compute_perievent(tsd, tref, windowsize=(1, 2, 3))
+    assert (
+        str(e_info.value)
+        == "windowsize should be a tuple of 2 numbers or a single number."
+    )
+
+    with pytest.raises(RuntimeError) as e_info:
+        nap.compute_perievent(tsd, tref, windowsize=(1, "2"))
+    assert (
+        str(e_info.value)
+        == "windowsize should be a tuple of 2 numbers or a single number."
+    )
+
 
 def test_compute_perievent_with_tsgroup():
     tsgroup = nap.TsGroup(
@@ -97,7 +123,9 @@ def test_compute_perievent_time_units():
     tsd = nap.Tsd(t=np.arange(100), d=np.arange(100))
     tref = nap.Ts(t=np.arange(10, 100, 10))
     for tu, fa in zip(["s", "ms", "us"], [1, 1e3, 1e6]):
-        peth = nap.compute_perievent(tsd, tref, windowsize=(-10 * fa, 10 * fa), time_unit=tu)
+        peth = nap.compute_perievent(
+            tsd, tref, windowsize=(-10 * fa, 10 * fa), time_unit=tu
+        )
         for i, j in zip(peth.keys(), np.arange(0, 100, 10)):
             np.testing.assert_array_almost_equal(peth[i].index, np.arange(-10, 10))
             np.testing.assert_array_almost_equal(peth[i].values, np.arange(j, j + 20))
@@ -127,7 +155,9 @@ def test_compute_perievent_continuous():
     np.testing.assert_array_almost_equal(
         pe.index.values, np.arange(windowsize[0], windowsize[-1] + 1)
     )
-    tmp = np.array([np.arange(t + windowsize[0], t + windowsize[1] + 1) for t in tref.t]).T
+    tmp = np.array(
+        [np.arange(t + windowsize[0], t + windowsize[1] + 1) for t in tref.t]
+    ).T
     np.testing.assert_array_almost_equal(pe.values, tmp)
 
     windowsize = (5, 10)
@@ -184,11 +214,15 @@ def test_compute_perievent_continuous_time_units():
     tref = nap.Ts(t=np.array([20, 60]))
     windowsize = (-5, 10)
     for tu, fa in zip(["s", "ms", "us"], [1, 1e3, 1e6]):
-        pe = nap.compute_perievent_continuous(tsd, tref, windowsize=(windowsize[0] * fa, windowsize[1] * fa), time_unit=tu)
+        pe = nap.compute_perievent_continuous(
+            tsd, tref, windowsize=(windowsize[0] * fa, windowsize[1] * fa), time_unit=tu
+        )
         np.testing.assert_array_almost_equal(
             pe.index.values, np.arange(windowsize[0], windowsize[1] + 1)
         )
-        tmp = np.array([np.arange(t + windowsize[0], t + windowsize[1] + 1) for t in tref.t]).T
+        tmp = np.array(
+            [np.arange(t + windowsize[0], t + windowsize[1] + 1) for t in tref.t]
+        ).T
         np.testing.assert_array_almost_equal(pe.values, tmp)
 
 
@@ -201,7 +235,10 @@ def test_compute_perievent_continuous_with_ep():
 
     assert pe.shape[1] == len(tref) - 1
     tmp = np.array(
-        [np.arange(t + windowsize[0], t + windowsize[1] + 1) for t in tref.restrict(ep).t]
+        [
+            np.arange(t + windowsize[0], t + windowsize[1] + 1)
+            for t in tref.restrict(ep).t
+        ]
     ).T
     np.testing.assert_array_almost_equal(pe.values, tmp)
 
@@ -237,26 +274,55 @@ def test_compute_perievent_continuous_with_ep():
 def test_compute_perievent_continuous_raise_error():
     tsd = nap.Tsd(t=np.arange(100), d=np.arange(100))
     tref = nap.Ts(t=np.arange(10, 100, 10))
-    with pytest.raises(AssertionError) as e_info:
-        nap.compute_perievent_continuous(tsd, [0, 1, 2], windowsize=(-10, 10))
-    assert str(e_info.value) == "tref should be a Ts or Tsd object."
+    with pytest.raises(TypeError) as e_info:
+        nap.compute_perievent_continuous(tsd, tref=[0, 1, 2], windowsize=(-10, 10))
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter tref must be of type ['Ts', 'Tsd', 'TsdFrame', 'TsdTensor']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_perievent_continuous([0, 1, 2], tref, windowsize=(-10, 10))
-    assert str(e_info.value) == "data should be a Tsd, TsdFrame or TsdTensor."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter timeseries must be of type ['Tsd', 'TsdFrame', 'TsdTensor']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_perievent_continuous(tsd, tref, windowsize={0: 1})
-    assert str(e_info.value) == "windowsize should be a tuple or int or float."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter windowsize must be of type ['tuple', 'Number']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_perievent_continuous(tsd, tref, windowsize=10, time_unit=1)
-    assert str(e_info.value) == "time_unit should be a str."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter time_unit must be of type ['str']."
+    )
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(RuntimeError) as e_info:
         nap.compute_perievent_continuous(tsd, tref, windowsize=10, time_unit="a")
     assert str(e_info.value) == "time_unit should be 's', 'ms' or 'us'"
 
-    with pytest.raises(AssertionError) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_perievent_continuous(tsd, tref, windowsize=10, ep="a")
-    assert str(e_info.value) == "ep should be an IntervalSet object."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter ep must be of type ['IntervalSet']."
+    )
+
+    with pytest.raises(RuntimeError) as e_info:
+        nap.compute_perievent_continuous(tsd, tref, windowsize=(1, 2, 3))
+    assert (
+        str(e_info.value)
+        == "windowsize should be a tuple of 2 numbers or a single number."
+    )
+
+    with pytest.raises(RuntimeError) as e_info:
+        nap.compute_perievent_continuous(tsd, tref, windowsize=(1, "2"))
+    assert (
+        str(e_info.value)
+        == "windowsize should be a tuple of 2 numbers or a single number."
+    )

tests/test_spike_trigger_average.py

@@ -33,6 +33,11 @@ def test_compute_spike_trigger_average_tsd():
     assert sta.shape == output.shape
     np.testing.assert_array_almost_equal(sta.values, output)
 
+    sta = nap.compute_event_trigger_average(spikes, feature, 0.2, (0.6, 0.6))
+    assert isinstance(sta, nap.TsdFrame)
+    assert sta.shape == output.shape
+    np.testing.assert_array_almost_equal(sta.values, output)
+
 
 def test_compute_spike_trigger_average_tsdframe():
     ep = nap.IntervalSet(0, 100)
@@ -149,41 +154,62 @@ def test_compute_spike_trigger_average_raise_error():
         {0: nap.Ts(t1), 1: nap.Ts(t1 - 0.1), 2: nap.Ts(t1 + 0.2)}, time_support=ep
     )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_event_trigger_average(feature, feature, 0.1, (0.5, 0.5), ep)
-    assert str(e_info.value) == "group should be a TsGroup."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter group must be of type ['TsGroup']."
+    )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_event_trigger_average(spikes, np.array(10), 0.1, (0.5, 0.5), ep)
-    assert str(e_info.value) == "Feature should be a Tsd, TsdFrame or TsdTensor"
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter feature must be of type ['Tsd', 'TsdFrame', 'TsdTensor']."
+    )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_event_trigger_average(spikes, feature, "0.1", (0.5, 0.5), ep)
-    assert str(e_info.value) == "binsize should be int or float."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter binsize must be of type ['Number']."
+    )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_event_trigger_average(
             spikes, feature, 0.1, (0.5, 0.5), ep, time_unit=1
         )
-    assert str(e_info.value) == "time_unit should be a str."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter time_unit must be of type ['str']."
+    )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(RuntimeError) as e_info:
         nap.compute_event_trigger_average(
             spikes, feature, 0.1, (0.5, 0.5), ep, time_unit="a"
         )
     assert str(e_info.value) == "time_unit should be 's', 'ms' or 'us'"
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(RuntimeError) as e_info:
         nap.compute_event_trigger_average(spikes, feature, 0.1, (0.5, 0.5, 0.5), ep)
-    assert str(e_info.value) == "windowsize should be a tuple of 2 elements (-t, +t)"
+    assert (
+        str(e_info.value)
+        == "windowsize should be a tuple of 2 numbers or a single number."
+    )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(RuntimeError) as e_info:
         nap.compute_event_trigger_average(spikes, feature, 0.1, ("a", "b"), ep)
-    assert str(e_info.value) == "windowsize should be a tuple of int/float"
+    assert (
+        str(e_info.value)
+        == "windowsize should be a tuple of 2 numbers or a single number."
+    )
 
-    with pytest.raises(Exception) as e_info:
+    with pytest.raises(TypeError) as e_info:
         nap.compute_event_trigger_average(spikes, feature, 0.1, (0.5, 0.5), [1, 2, 3])
-    assert str(e_info.value) == "ep should be an IntervalSet object."
+    assert (
+        str(e_info.value)
+        == "Invalid type. Parameter ep must be of type ['IntervalSet']."
+    )
 
 
 def test_compute_spike_trigger_average_time_unit():