Working version of build_tensor

Author: gviejo

pynapple/process/warping.py

@@ -25,6 +25,7 @@ def wrapper(*args, **kwargs):
             "time_unit": (str,),
             "align": (str,),
             "padding_value": (Number,),
+            "num_bin" : (int,)
         }
         for param, param_type in parameters_type.items():
             if param in kwargs:
@@ -39,25 +40,35 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def _build_tensor_from_tsgroup(input, ep, binsize, align, padding_value):
+def _build_tensor_from_tsgroup(input, ep, binsize, align, padding_value, time_unit):
     # Determine size of tensor
+    binsize = float(nap.TsIndex.format_timestamps(np.array([binsize]), time_unit)[0])
     n_t = int(np.max(np.ceil((ep.end + binsize - ep.start) / binsize)))
-    output = np.ones(shape=(len(input), len(ep), n_t)) * padding_value
     count = input.count(bin_size=binsize, ep=ep)
 
+    if len(count.shape) == 1:
+        output = np.ones(shape=(1, len(ep), n_t)) * padding_value
+    else:
+        output = np.ones(shape=(count.shape[1], len(ep), n_t)) * padding_value
+
+    n_ep = np.zeros(len(ep), dtype="int")  # To trim to the minimum length
+
     if align == "start":
         for i in range(len(ep)):
             tmp = count.get(ep.start[i], ep.end[i]).values
+            n_ep[i] = tmp.shape[0]
             output[:, i, 0 : tmp.shape[0]] = np.transpose(tmp)
-        if np.all(np.isnan(output[:, :, -1])):
-            output = output[:, :, 0:-1]
+        output = output[:, :, 0 : np.max(n_ep)]
 
     if align == "end":
         for i in range(len(ep)):
             tmp = count.get(ep.start[i], ep.end[i]).values
+            n_ep[i] = tmp.shape[0]
             output[:, i, -tmp.shape[0] :] = np.transpose(tmp)
-        if np.all(np.isnan(output[:, :, 0])):
-            output = output[:, :, 1:]
+        output = output[:, :, -np.max(n_ep) :]
+
+    if len(count.shape) == 1:  # Removing first axis if Ts.
+        output = output[0]
 
     return output
 
@@ -100,7 +111,7 @@ def build_tensor(
 
     Parameters
     ----------
-    input : Tsd, TsdFrame, TsdTensor or TsGroup
+    input : Ts, Tsd, TsdFrame, TsdTensor or TsGroup
         Input to slice and align to the trials within the `ep` parameter.
     ep : IntervalSet
         Epochs holding the trials. Each interval can be of unequal size.
@@ -121,65 +132,96 @@ def build_tensor(
     RuntimeError
         If `time_unit` not in ["s", "ms", "us"]
 
-
     Examples
     --------
-
-
+    >>> import pynapple as nap
+    >>> import numpy as np
+    >>> group = nap.TsGroup({0:nap.Ts(t=np.arange(0, 100))})
+    >>> ep = nap.IntervalSet(start=np.arange(20, 100, 20), end=np.arange(20, 100, 20) + np.arange(2, 10, 2))
+    >>> print(ep)
+      index    start    end
+          0       20     22
+          1       40     44
+          2       60     66
+          3       80     88
+    shape: (4, 2), time unit: sec.
+
+    Create a trial-based tensor by counting events within 1 second bin for each interval of `ep`.
+
+    >>> tensor = nap.build_tensor(group, ep, binsize=1)
+    >>> tensor
+    array([[[ 1.,  1., nan, nan, nan, nan, nan, nan],
+            [ 1.,  1.,  1.,  1., nan, nan, nan, nan],
+            [ 1.,  1.,  1.,  1.,  1.,  1., nan, nan],
+            [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.]]])
+
+    By default, the time series are aligned to the start of the epochs. The parameter `align` control this behavior.
+
+    >>> tensor = nap.build_tensor(group, ep, binsize=1, align="end")
+    >>> tensor
+    array([[[nan, nan, nan, nan, nan, nan,  1.,  1.],
+            [nan, nan, nan, nan,  1.,  1.,  1.,  1.],
+            [nan, nan,  1.,  1.,  1.,  1.,  1.,  1.],
+            [ 1.,  1.,  1.,  1.,  1.,  1.,  1.,  1.]]])
+
+    This function works for any time series.
+
+    >>> tsdframe = nap.TsdFrame(t=np.arange(100), d=np.arange(200).reshape(2,100).T)
+    >>> tensor = nap.build_tensor(tsdframe, ep)
+    >>> tensor
+    array([[[ 20.,  21.,  22.,  nan,  nan,  nan,  nan,  nan,  nan],
+            [ 40.,  41.,  42.,  43.,  44.,  nan,  nan,  nan,  nan],
+            [ 60.,  61.,  62.,  63.,  64.,  65.,  66.,  nan,  nan],
+            [ 80.,  81.,  82.,  83.,  84.,  85.,  86.,  87.,  88.]],
+           [[120., 121., 122.,  nan,  nan,  nan,  nan,  nan,  nan],
+            [140., 141., 142., 143., 144.,  nan,  nan,  nan,  nan],
+            [160., 161., 162., 163., 164., 165., 166.,  nan,  nan],
+            [180., 181., 182., 183., 184., 185., 186., 187., 188.]]])
 
     """
     if time_unit not in ["s", "ms", "us"]:
         raise RuntimeError("time_unit should be 's', 'ms' or 'us'")
     if align not in ["start", "end"]:
         raise RuntimeError("align should be 'start' or 'end'")
 
-    if isinstance(input, nap.TsGroup):
+    if isinstance(input, (nap.TsGroup, nap.Ts)):
         if not isinstance(binsize, Number):
-            raise RuntimeError("When input is a TsGroup, binsize should be specified")
-        return _build_tensor_from_tsgroup(input, ep, binsize, align, padding_value)
-
-    if isinstance(input, (nap.Tsd, nap.TsdFrame, nap.TsdTensor)):
+            raise RuntimeError(
+                "When input is a TsGroup or Ts object, binsize should be specified"
+            )
+        return _build_tensor_from_tsgroup(
+            input, ep, binsize, align, padding_value, time_unit
+        )
+    else:
         return _build_tensor_from_tsd(input, ep, align, padding_value)
 
 
 @_validate_warping_inputs
-def warp_tensor(input, ep, num_bin=None, align="start"):
+def warp_tensor(input, ep, num_bin=None):
     """
-    Return time-warped trial-based tensor from an IntervalSet object.
+    Return linearly time-warped trial-based tensor from an IntervalSet object.
 
-    - If `input` is a `TsGroup`, returns a numpy array of shape (number of group element, number of trial, number of time bins). The `binsize` parameter determines the number of time bins.
+    - If `input` is a `TsGroup`, returns a numpy array of shape (number of group element, number of trial, `num_bin`).
 
-    - If `input` is `Tsd`, `TsdFrame` or `TsdTensor`, returns a numpy array of shape (shape of time series, number of trial, number of  time points).
+    - If `input` is `Tsd`, `TsdFrame` or `TsdTensor`, returns a numpy array of shape (shape of time series, number of trial, `num_bin`).
 
 
     Parameters
     ----------
-    input : Tsd, TsdFrame, TsdTensor or TsGroup
-        Returns a numpy array.
+    input : Ts, Tsd, TsdFrame, TsdTensor or TsGroup
+        Input object
     ep : IntervalSet
         Epochs holding the trials. Each interval can be of unequal size.
-    binsize : Number, optional
-    align: str, optional
-        How to align the time series ('start' [default], 'end')
-    padding_value: Number, optional
-        How to pad the array if unequal intervals. Default is np.nan.
-    time_unit : str, optional
-        Time units of the binsize parameter ('s' [default], 'ms', 'us').
+    num_bin : int
 
     Returns
     -------
     numpy.ndarray
 
-    Raises
-    ------
-    RuntimeError
-        If `time_unit` not in ["s", "ms", "us"]
-
-
     Examples
     --------
 
 
 
     """
-    pass
+

tests/test_warping.py

@@ -103,7 +103,8 @@ def test_build_tensor_runtime_error():
     )
 
     with pytest.raises(
-        RuntimeError, match=r"When input is a TsGroup, binsize should be specified"
+        RuntimeError,
+        match=r"When input is a TsGroup or Ts object, binsize should be specified",
     ):
         nap.build_tensor(group, ep)
 
@@ -140,3 +141,151 @@ def test_build_tensor_with_group():
 
     tensor = nap.build_tensor(group, ep, binsize=1, align="end")
     np.testing.assert_array_almost_equal(tensor, np.flip(expected, axis=2))
+
+    tensor = nap.build_tensor(group, ep, binsize=1, time_unit="s")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(group, ep, binsize=1e3, time_unit="ms")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(group, ep, binsize=1e6, time_unit="us")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(group, ep, binsize=1, align="start", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+
+def test_build_tensor_with_ts():
+    ts = nap.Ts(t=np.arange(0, 100))
+    ep = nap.IntervalSet(
+        start=np.arange(0, 100, 20), end=np.arange(0, 100, 20) + np.arange(0, 10, 2)
+    )
+
+    expected = np.ones((len(ep), 8)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, 0:k] = 1
+
+    tensor = nap.build_tensor(ts, ep, binsize=1)
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(ts, ep, binsize=1, align="start")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(ts, ep, binsize=1, align="end")
+    np.testing.assert_array_almost_equal(tensor, np.flip(expected, axis=1))
+
+    tensor = nap.build_tensor(ts, ep, binsize=1, time_unit="s")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(ts, ep, binsize=1e3, time_unit="ms")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(ts, ep, binsize=1e6, time_unit="us")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(ts, ep, binsize=1, align="start", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+
+def test_build_tensor_with_tsd():
+    tsd = nap.Tsd(t=np.arange(100), d=np.arange(100))
+    ep = nap.IntervalSet(
+        start=np.arange(0, 100, 20), end=np.arange(0, 100, 20) + np.arange(0, 10, 2)
+    )
+
+    expected = np.ones((len(ep), 9)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, 0 : k + 1] = np.arange(k * 10, k * 10 + k + 1)
+
+    tensor = nap.build_tensor(tsd, ep)
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsd, ep, align="start")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsd, ep, align="start", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    expected = np.ones((len(ep), 9)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, -k - 1 :] = np.arange(k * 10, k * 10 + k + 1)
+
+    tensor = nap.build_tensor(tsd, ep, align="end")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsd, ep, align="end", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+
+def test_build_tensor_with_tsdframe():
+    tsdframe = nap.TsdFrame(t=np.arange(100), d=np.tile(np.arange(100)[:, None], 3))
+    ep = nap.IntervalSet(
+        start=np.arange(0, 100, 20), end=np.arange(0, 100, 20) + np.arange(0, 10, 2)
+    )
+
+    expected = np.ones((len(ep), 9)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, 0 : k + 1] = np.arange(k * 10, k * 10 + k + 1)
+    expected = np.repeat(expected[None, :], 3, axis=0)
+
+    tensor = nap.build_tensor(tsdframe, ep)
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsdframe, ep, align="start")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsdframe, ep, align="start", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    expected = np.ones((len(ep), 9)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, -k - 1 :] = np.arange(k * 10, k * 10 + k + 1)
+    expected = np.repeat(expected[None, :], 3, axis=0)
+
+    tensor = nap.build_tensor(tsdframe, ep, align="end")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsdframe, ep, align="end", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+
+def test_build_tensor_with_tsdtensor():
+    tsdtensor = nap.TsdTensor(
+        t=np.arange(100), d=np.tile(np.arange(100)[:, None, None], (1, 2, 3))
+    )
+    ep = nap.IntervalSet(
+        start=np.arange(0, 100, 20), end=np.arange(0, 100, 20) + np.arange(0, 10, 2)
+    )
+
+    expected = np.ones((len(ep), 9)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, 0 : k + 1] = np.arange(k * 10, k * 10 + k + 1)
+    expected = np.tile(expected[None, None, :, :], (2, 3, 1, 1))
+
+    tensor = nap.build_tensor(tsdtensor, ep)
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsdtensor, ep, align="start")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsdtensor, ep, align="start", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    expected = np.ones((len(ep), 9)) * np.nan
+    for i, k in zip(range(len(ep)), range(2, 10, 2)):
+        expected[i, -k - 1 :] = np.arange(k * 10, k * 10 + k + 1)
+    expected = np.tile(expected[None, None, :, :], (2, 3, 1, 1))
+
+    tensor = nap.build_tensor(tsdtensor, ep, align="end")
+    np.testing.assert_array_almost_equal(tensor, expected)
+
+    tensor = nap.build_tensor(tsdtensor, ep, align="end", padding_value=-1)
+    expected[np.isnan(expected)] = -1
+    np.testing.assert_array_almost_equal(tensor, expected)