Update test_time_delay.py

1 files changed, 18 insertions, 19 deletions

diff --git a/src/python/test/test_time_delay.py b/src/python/test/test_time_delay.py
index 7b6562a5..f652fc88 100755
--- a/src/python/test/test_time_delay.py
+++ b/src/python/test/test_time_delay.py
@@ -7,31 +7,30 @@ def test_normal():
     # Normal case.
     prep = TimeDelayEmbedding()
     attractor = prep(ts)
-    assert (attractor[0] == np.array([1, 2, 3])
-            print(attractor[0].all()))
-    assert (attractor[1] == np.array([2, 3, 4]))
-    assert (attractor[2] == np.array([3, 4, 5]))
-    assert (attractor[3] == np.array([4, 5, 6]))
-    assert (attractor[4] == np.array([5, 6, 7]))
-    assert (attractor[5] == np.array([6, 7, 8]))
-    assert (attractor[6] == np.array([7, 8, 9]))
-    assert (attractor[7] == np.array([8, 9, 10]))
+    assert (attractor[0].all() == np.array([1, 2, 3]))
+    assert (attractor[1].all() == np.array([2, 3, 4]))
+    assert (attractor[2].all() == np.array([3, 4, 5]))
+    assert (attractor[3].all() == np.array([4, 5, 6]))
+    assert (attractor[4].all() == np.array([5, 6, 7]))
+    assert (attractor[5].all() == np.array([6, 7, 8]))
+    assert (attractor[6].all() == np.array([7, 8, 9]))
+    assert (attractor[7].all() == np.array([8, 9, 10]))
     # Delay = 3
     prep = TimeDelayEmbedding(delay=3)
     attractor = prep(ts)
-    assert (attractor[0] == np.array([1, 4, 7]))
-    assert (attractor[1] == np.array([2, 5, 8]))
-    assert (attractor[2] == np.array([3, 6, 9]))
-    assert (attractor[3] == np.array([4, 7, 10]))
+    assert (attractor[0].all() == np.array([1, 4, 7]))
+    assert (attractor[1].all() == np.array([2, 5, 8]))
+    assert (attractor[2].all() == np.array([3, 6, 9]))
+    assert (attractor[3].all() == np.array([4, 7, 10]))
     # Skip = 3
     prep = TimeDelayEmbedding(skip=3)
     attractor = prep(ts)
-    assert (attractor[0] == np.array([1, 2, 3]))
-    assert (attractor[1] == np.array([4, 5, 6]))
-    assert (attractor[2] == np.array([7, 8, 9]))
+    assert (attractor[0].all() == np.array([1, 2, 3]))
+    assert (attractor[1].all() == np.array([4, 5, 6]))
+    assert (attractor[2].all() == np.array([7, 8, 9]))
     # Delay = 2 / Skip = 2
     prep = TimeDelayEmbedding(delay=2, skip=2)
     attractor = prep(ts)
-    assert (attractor[0] == np.array([1, 3, 5]))
-    assert (attractor[1] == np.array([3, 5, 7]))
-    assert (attractor[2] == np.array([5, 7, 9]))
+    assert (attractor[0].all() == np.array([1, 3, 5]))
+    assert (attractor[1].all() == np.array([3, 5, 7]))
+    assert (attractor[2].all() == np.array([5, 7, 9]))