diff --git a/src/pytrxmat/trx.py b/src/pytrxmat/trx.py index 8c2e01c42b75725d55ad86289142fa2c5352870b..597f8847a54656bf0e6d24be5b1f236d7d2eb5c4 100644 --- a/src/pytrxmat/trx.py +++ b/src/pytrxmat/trx.py @@ -107,6 +107,7 @@ class TRX: * tuple(float, float) : return all times t_sel such that t[0] <= t_sel < t[1]. * tuple(float, int) : return up to t[1] points after the intial time t[0]. * tuple(int, float) : return up to t[0] points before the finalal time t[1]. + * tuple(int, int) : slice the time axis along t[0]:t[1]. There is no guarantee that the shapes are the same across larvae due to missing data. If no points verify the condition, returns an array whith second dimension equal to 0. @@ -159,8 +160,10 @@ class TRX: arrays = [a[:, s:s+t2] if s is not None else np.empty((a.shape[0],0)) for a, s in zip(arrays, start_points)] elif isinstance(t1, int) and isinstance(t2, float): raise NotImplementedError() + elif isinstance(t1, int) and isinstance(t2, int): + arrays = [a[:, t1:t2] for a in arrays] else: - raise ValueError("'t' can only be of type Tuple(float, float), Tuple(float, int), or Tuple(int, float)") + raise ValueError("'t' can only be of type Tuple(float, float), Tuple(float, int), Tuple(int, int), or Tuple(int, float)") if 'return_columns' in kwargs and kwargs['return_columns']: return arrays, columns