def test_shift_time_indices(): A = (1 + np.arange(2 * 4 * 6)).reshape((2, 4, 6)) B = convolve.shift_time_indices(A, trim=False) C = np.zeros_like(A) for i, j, k in product(*map(xrange, A.shape)): if (i + j) < A.shape[1] and (i + j + k) < A.shape[2]: C[i, j, k] = A[i, i + j, i + j + k] assert_equal(B, C)